YARN-2753. Fixed a bunch of bugs in the NodeLabelsManager classes. Contributed by Zhihai xu.

(cherry picked from commit 4cfd5bc7c18bb9a828f573b5c4d2b13fa28e732a)

hadoop-yarn-project/CHANGES.txt

@@ -831,6 +831,9 @@ Release 2.6.0 - UNRELEASED
     YARN-2803. MR distributed cache not working correctly on Windows after
     NodeManager privileged account changes. (Craig Welch via cnauroth)
 
+    YARN-2753. Fixed a bunch of bugs in the NodeLabelsManager classes. (Zhihai xu
+    via vinodkv)
+
 Release 2.5.2 - UNRELEASED
 
   INCOMPATIBLE CHANGES

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java

@@ -144,9 +144,7 @@ public class CommonNodeLabelsManager extends AbstractService {
 
     @Override
     public void handle(NodeLabelsStoreEvent event) {
-      if (isInState(STATE.STARTED)) {
-        handleStoreEvent(event);
-      }
+      handleStoreEvent(event);
     }
   }
   
@@ -256,7 +254,7 @@ public class CommonNodeLabelsManager extends AbstractService {
     if (null == labels || labels.isEmpty()) {
       return;
     }
-
+    Set<String> newLabels = new HashSet<String>();
     labels = normalizeLabels(labels);
 
     // do a check before actual adding them, will throw exception if any of them
@@ -266,11 +264,15 @@ public class CommonNodeLabelsManager extends AbstractService {
     }
 
     for (String label : labels) {
-      this.labelCollections.put(label, new Label());
+      // shouldn't overwrite it to avoid changing the Label.resource
+      if (this.labelCollections.get(label) == null) {
+        this.labelCollections.put(label, new Label());
+        newLabels.add(label);
+      }
     }
-    if (null != dispatcher) {
+    if (null != dispatcher && !newLabels.isEmpty()) {
       dispatcher.getEventHandler().handle(
-          new StoreNewClusterNodeLabels(labels));
+          new StoreNewClusterNodeLabels(newLabels));
     }
 
     LOG.info("Add labels: [" + StringUtils.join(labels.iterator(), ",") + "]");
@@ -453,12 +455,15 @@ public class CommonNodeLabelsManager extends AbstractService {
         LOG.error(msg);
         throw new IOException(msg);
       }
-      
-      if (labels == null || labels.isEmpty()) {
+
+      // the labels will never be null
+      if (labels.isEmpty()) {
         continue;
       }
 
-      if (!originalLabels.containsAll(labels)) {
+      // originalLabels may be null,
+      // because when a Node is created, Node.labels can be null.
+      if (originalLabels == null || !originalLabels.containsAll(labels)) {
         String msg =
             "Try to remove labels = [" + StringUtils.join(labels, ",")
                 + "], but not all labels contained by NM=" + nodeId;

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/TestRMNodeLabelsManager.java

@@ -74,6 +74,13 @@ public class TestRMNodeLabelsManager extends NodeLabelTestBase {
     Assert.assertEquals(mgr.getResourceByLabel("p1", null),
         Resources.add(SMALL_RESOURCE, LARGE_NODE));
 
+    // check add labels multiple times shouldn't overwrite
+    // original attributes on labels like resource
+    mgr.addToCluserNodeLabels(toSet("p1", "p4"));
+    Assert.assertEquals(mgr.getResourceByLabel("p1", null),
+        Resources.add(SMALL_RESOURCE, LARGE_NODE));
+    Assert.assertEquals(mgr.getResourceByLabel("p4", null), EMPTY_RESOURCE);
+
     // change the large NM to small, check if resource updated
     mgr.updateNodeResource(NodeId.newInstance("n1", 2), SMALL_RESOURCE);
     Assert.assertEquals(mgr.getResourceByLabel("p1", null),
@@ -374,7 +381,7 @@ public class TestRMNodeLabelsManager extends NodeLabelTestBase {
     Assert.assertEquals(clusterResource,
         mgr.getQueueResource("Q5", q5Label, clusterResource));
   }
-  
+
   @Test(timeout=5000)
   public void testGetLabelResourceWhenMultipleNMsExistingInSameHost() throws IOException {
     // active two NM to n1, one large and one small
@@ -401,4 +408,24 @@ public class TestRMNodeLabelsManager extends NodeLabelTestBase {
             mgr.getResourceByLabel("p1", null),
             Resources.multiply(SMALL_RESOURCE, 2));
   }
+
+  @Test(timeout = 5000)
+  public void testRemoveLabelsFromNode() throws Exception {
+    mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), toSet("p1"),
+            toNodeId("n2"), toSet("p2"), toNodeId("n3"), toSet("p3")));
+    // active one NM to n1:1
+    mgr.activateNode(NodeId.newInstance("n1", 1), SMALL_RESOURCE);
+    try {
+      mgr.removeLabelsFromNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p1")));
+      Assert.fail("removeLabelsFromNode should trigger IOException");
+    } catch (IOException e) {
+    }
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p1")));
+    try {
+      mgr.removeLabelsFromNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p1")));
+    } catch (IOException e) {
+      Assert.fail("IOException from removeLabelsFromNode " + e);
+    }
+  }
 }