From f9139ac8f60184a82a8bb315237bea04bdb98ec8 Mon Sep 17 00:00:00 2001
From: Jian He <jianhe@apache.org>
Date: Wed, 2 Aug 2017 00:03:25 -0700
Subject: [PATCH] YARN-6872. [Addendum patch] Ensure apps could run given
 NodeLabels are disabled post RM switchover/restart. Contributed by Sunil G

---
 .../scheduler/AbstractYarnScheduler.java       | 18 ++++--------------
 .../scheduler/AppSchedulingInfo.java           |  6 +++---
 .../scheduler/SchedulerApplicationAttempt.java |  2 +-
 3 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AbstractYarnScheduler.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AbstractYarnScheduler.java
index b40ce407618..d506f4d2cce 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AbstractYarnScheduler.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AbstractYarnScheduler.java
@@ -519,20 +519,10 @@ public abstract class AbstractYarnScheduler
     container.setVersion(status.getVersion());
     ApplicationAttemptId attemptId =
         container.getId().getApplicationAttemptId();
-    String labelExpression = status.getNodeLabelExpression();
-    // If NodeLabel is disabled but recovered container has label expression
-    // its better to suppress that and considered as default label.
-    if (!status.getNodeLabelExpression().isEmpty() && !YarnConfiguration
-        .areNodeLabelsEnabled(rmContext.getYarnConfiguration())) {
-      labelExpression = RMNodeLabelsManager.NO_LABEL;
-    }
-
-    RMContainer rmContainer =
-        new RMContainerImpl(container,
-            SchedulerRequestKey.extractFrom(container), attemptId,
-            node.getNodeID(), applications.get(
-            attemptId.getApplicationId()).getUser(), rmContext,
-            status.getCreationTime(), labelExpression);
+    RMContainer rmContainer = new RMContainerImpl(container,
+        SchedulerRequestKey.extractFrom(container), attemptId, node.getNodeID(),
+        applications.get(attemptId.getApplicationId()).getUser(), rmContext,
+        status.getCreationTime(), status.getNodeLabelExpression());
     return rmContainer;
   }
 
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
index 851c9f5587b..8acf7d55dbb 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
@@ -524,7 +524,7 @@ public class AppSchedulingInfo {
     this.placesBlacklistedByApp = appInfo.getBlackList();
   }
 
-  public void recoverContainer(RMContainer rmContainer) {
+  public void recoverContainer(RMContainer rmContainer, String partition) {
     try {
       this.writeLock.lock();
       QueueMetrics metrics = queue.getMetrics();
@@ -540,8 +540,8 @@ public class AppSchedulingInfo {
         return;
       }
 
-      metrics.allocateResources(rmContainer.getNodeLabelExpression(),
-          user, 1, rmContainer.getAllocatedResource(), false);
+      metrics.allocateResources(partition, user, 1,
+          rmContainer.getAllocatedResource(), false);
     } finally {
       this.writeLock.unlock();
     }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/SchedulerApplicationAttempt.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/SchedulerApplicationAttempt.java
index 4b0bf91e250..8b2f9db927b 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/SchedulerApplicationAttempt.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/SchedulerApplicationAttempt.java
@@ -1103,7 +1103,7 @@ public class SchedulerApplicationAttempt implements SchedulableEntity {
     try {
       writeLock.lock();
       // recover app scheduling info
-      appSchedulingInfo.recoverContainer(rmContainer);
+      appSchedulingInfo.recoverContainer(rmContainer, node.getPartition());
 
       if (rmContainer.getState().equals(RMContainerState.COMPLETED)) {
         return;