소스 검색

YARN-3075. NodeLabelsManager implementation to retrieve label to node mapping (Varun Saxena via wangda)

(cherry picked from commit 5bd984691b3e3df5d8aac910295138996e7a5839)
Wangda Tan 10 년 전
부모
커밋
205e15c4a3

+ 3 - 0
hadoop-yarn-project/CHANGES.txt

@@ -200,6 +200,9 @@ Release 2.7.0 - UNRELEASED
     YARN-3098. Created common QueueCapacities class in Capacity Scheduler to
     YARN-3098. Created common QueueCapacities class in Capacity Scheduler to
     track capacities-by-labels of queues. (Wangda Tan via jianhe)
     track capacities-by-labels of queues. (Wangda Tan via jianhe)
 
 
+    YARN-3075. NodeLabelsManager implementation to retrieve label to node 
+    mapping (Varun Saxena via wangda)
+
   OPTIMIZATIONS
   OPTIMIZATIONS
 
 
   BUG FIXES
   BUG FIXES

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

@@ -121,15 +121,17 @@ public class CommonNodeLabelsManager extends AbstractService {
     public Set<String> labels;
     public Set<String> labels;
     public Resource resource;
     public Resource resource;
     public boolean running;
     public boolean running;
+    public NodeId nodeId;
     
     
-    protected Node() {
+    protected Node(NodeId nodeid) {
       labels = null;
       labels = null;
       resource = Resource.newInstance(0, 0);
       resource = Resource.newInstance(0, 0);
       running = false;
       running = false;
+      nodeId = nodeid;
     }
     }
     
     
     public Node copy() {
     public Node copy() {
-      Node c = new Node();
+      Node c = new Node(nodeId);
       if (labels != null) {
       if (labels != null) {
         c.labels =
         c.labels =
             Collections.newSetFromMap(new ConcurrentHashMap<String, Boolean>());
             Collections.newSetFromMap(new ConcurrentHashMap<String, Boolean>());
@@ -465,7 +467,27 @@ public class CommonNodeLabelsManager extends AbstractService {
       }
       }
     }
     }
   }
   }
-  
+
+  private void addNodeToLabels(NodeId node, Set<String> labels) {
+    for(String l : labels) {
+      labelCollections.get(l).addNodeId(node);
+    }
+  }
+
+  private void removeNodeFromLabels(NodeId node, Set<String> labels) {
+    for(String l : labels) {
+      labelCollections.get(l).removeNodeId(node);
+    }
+  }
+
+  private void replaceNodeForLabels(NodeId node, Set<String> oldLabels,
+      Set<String> newLabels) {
+    if(oldLabels != null) {
+      removeNodeFromLabels(node, oldLabels);
+    }
+    addNodeToLabels(node, newLabels);
+  }
+
   @SuppressWarnings("unchecked")
   @SuppressWarnings("unchecked")
   protected void internalUpdateLabelsOnNodes(
   protected void internalUpdateLabelsOnNodes(
       Map<NodeId, Set<String>> nodeToLabels, NodeLabelUpdateOperation op)
       Map<NodeId, Set<String>> nodeToLabels, NodeLabelUpdateOperation op)
@@ -473,6 +495,7 @@ public class CommonNodeLabelsManager extends AbstractService {
     // do update labels from nodes
     // do update labels from nodes
     Map<NodeId, Set<String>> newNMToLabels =
     Map<NodeId, Set<String>> newNMToLabels =
         new HashMap<NodeId, Set<String>>();
         new HashMap<NodeId, Set<String>>();
+    Set<String> oldLabels;
     for (Entry<NodeId, Set<String>> entry : nodeToLabels.entrySet()) {
     for (Entry<NodeId, Set<String>> entry : nodeToLabels.entrySet()) {
       NodeId nodeId = entry.getKey();
       NodeId nodeId = entry.getKey();
       Set<String> labels = entry.getValue();
       Set<String> labels = entry.getValue();
@@ -481,26 +504,32 @@ public class CommonNodeLabelsManager extends AbstractService {
       if (nodeId.getPort() == WILDCARD_PORT) {
       if (nodeId.getPort() == WILDCARD_PORT) {
         Host host = nodeCollections.get(nodeId.getHost());
         Host host = nodeCollections.get(nodeId.getHost());
         switch (op) {
         switch (op) {
-        case REMOVE: 
+        case REMOVE:
+          removeNodeFromLabels(nodeId, labels);
           host.labels.removeAll(labels);
           host.labels.removeAll(labels);
           for (Node node : host.nms.values()) {
           for (Node node : host.nms.values()) {
             if (node.labels != null) {
             if (node.labels != null) {
               node.labels.removeAll(labels);
               node.labels.removeAll(labels);
             }
             }
+            removeNodeFromLabels(node.nodeId, labels);
           }
           }
           break;
           break;
         case ADD:
         case ADD:
+          addNodeToLabels(nodeId, labels);
           host.labels.addAll(labels);
           host.labels.addAll(labels);
           for (Node node : host.nms.values()) {
           for (Node node : host.nms.values()) {
             if (node.labels != null) {
             if (node.labels != null) {
               node.labels.addAll(labels);
               node.labels.addAll(labels);
             }
             }
+            addNodeToLabels(node.nodeId, labels);
           }
           }
           break;
           break;
         case REPLACE:
         case REPLACE:
+          replaceNodeForLabels(nodeId, host.labels, labels);
           host.labels.clear();
           host.labels.clear();
           host.labels.addAll(labels);
           host.labels.addAll(labels);
           for (Node node : host.nms.values()) {
           for (Node node : host.nms.values()) {
+            replaceNodeForLabels(node.nodeId, node.labels, labels);
             node.labels = null;
             node.labels = null;
           }
           }
           break;
           break;
@@ -514,14 +543,20 @@ public class CommonNodeLabelsManager extends AbstractService {
           // Add and replace
           // Add and replace
           createNodeIfNonExisted(nodeId);
           createNodeIfNonExisted(nodeId);
           Node nm = getNMInNodeSet(nodeId);
           Node nm = getNMInNodeSet(nodeId);
-          if (nm.labels == null) {
-            nm.labels = new HashSet<String>();
-          }
           switch (op) {
           switch (op) {
           case ADD:
           case ADD:
+            addNodeToLabels(nodeId, labels);
+            if (nm.labels == null) { 
+              nm.labels = new HashSet<String>();
+            }
             nm.labels.addAll(labels);
             nm.labels.addAll(labels);
             break;
             break;
           case REPLACE:
           case REPLACE:
+            oldLabels = getLabelsByNode(nodeId);
+            replaceNodeForLabels(nodeId, oldLabels, labels);
+            if (nm.labels == null) { 
+              nm.labels = new HashSet<String>();
+            }
             nm.labels.clear();
             nm.labels.clear();
             nm.labels.addAll(labels);
             nm.labels.addAll(labels);
             break;
             break;
@@ -531,6 +566,7 @@ public class CommonNodeLabelsManager extends AbstractService {
           newNMToLabels.put(nodeId, nm.labels);
           newNMToLabels.put(nodeId, nm.labels);
         } else {
         } else {
           // remove
           // remove
+          removeNodeFromLabels(nodeId, labels);
           Node nm = getNMInNodeSet(nodeId);
           Node nm = getNMInNodeSet(nodeId);
           if (nm.labels != null) {
           if (nm.labels != null) {
             nm.labels.removeAll(labels);
             nm.labels.removeAll(labels);
@@ -646,6 +682,52 @@ public class CommonNodeLabelsManager extends AbstractService {
     }
     }
   }
   }
 
 
+  /**
+   * Get mapping of labels to nodes for all the labels.
+   *
+   * @return labels to nodes map
+   */
+  public Map<String, Set<NodeId>> getLabelsToNodes() {
+    try {
+      readLock.lock();
+      return getLabelsToNodes(labelCollections.keySet());
+    } finally {
+      readLock.unlock();
+    }
+  }
+
+  /**
+   * Get mapping of labels to nodes for specified set of labels.
+   *
+   * @param labels set of labels for which labels to nodes mapping will be
+   *        returned.
+   * @return labels to nodes map
+   */
+  public Map<String, Set<NodeId>> getLabelsToNodes(Set<String> labels) {
+    try {
+      readLock.lock();
+      Map<String, Set<NodeId>> labelsToNodes =
+          new HashMap<String, Set<NodeId>>();
+      for (String label : labels) {
+        if(label.equals(NO_LABEL)) {
+          continue;
+        }
+        NodeLabel nodeLabelInfo = labelCollections.get(label);
+        if(nodeLabelInfo != null) {
+          Set<NodeId> nodeIds = nodeLabelInfo.getAssociatedNodeIds();
+          if (!nodeIds.isEmpty()) {
+            labelsToNodes.put(label, nodeIds);
+          }
+        } else {
+          LOG.warn("getLabelsToNodes : Label [" + label + "] cannot be found");
+        }
+      }      
+      return Collections.unmodifiableMap(labelsToNodes);
+    } finally {
+      readLock.unlock();
+    }
+  }
+
   /**
   /**
    * Get existing valid labels in repository
    * Get existing valid labels in repository
    * 
    * 
@@ -741,7 +823,7 @@ public class CommonNodeLabelsManager extends AbstractService {
     }
     }
     Node nm = host.nms.get(nodeId);
     Node nm = host.nms.get(nodeId);
     if (null == nm) {
     if (null == nm) {
-      host.nms.put(nodeId, new Node());
+      host.nms.put(nodeId, new Node(nodeId));
     }
     }
   }
   }
   
   

+ 17 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeLabel.java

@@ -18,7 +18,10 @@
 
 
 package org.apache.hadoop.yarn.nodelabels;
 package org.apache.hadoop.yarn.nodelabels;
 
 
+import java.util.HashSet;
+import java.util.Set;
 import org.apache.commons.lang.StringUtils;
 import org.apache.commons.lang.StringUtils;
+import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.util.resource.Resources;
 import org.apache.hadoop.yarn.util.resource.Resources;
 
 
@@ -26,6 +29,7 @@ public class NodeLabel implements Comparable<NodeLabel> {
   private Resource resource;
   private Resource resource;
   private int numActiveNMs;
   private int numActiveNMs;
   private String labelName;
   private String labelName;
+  private Set<NodeId> nodeIds;
 
 
   public NodeLabel(String labelName) {
   public NodeLabel(String labelName) {
     this(labelName, Resource.newInstance(0, 0), 0);
     this(labelName, Resource.newInstance(0, 0), 0);
@@ -35,8 +39,21 @@ public class NodeLabel implements Comparable<NodeLabel> {
     this.labelName = labelName;
     this.labelName = labelName;
     this.resource = res;
     this.resource = res;
     this.numActiveNMs = activeNMs;
     this.numActiveNMs = activeNMs;
+    this.nodeIds = new HashSet<NodeId>();
+  }
+
+  public void addNodeId(NodeId node) {
+    nodeIds.add(node);
+  }
+
+  public void removeNodeId(NodeId node) {
+    nodeIds.remove(node);
   }
   }
   
   
+  public Set<NodeId> getAssociatedNodeIds() {
+    return new HashSet<NodeId>(nodeIds);
+  }
+
   public void addNode(Resource nodeRes) {
   public void addNode(Resource nodeRes) {
     Resources.addTo(resource, nodeRes);
     Resources.addTo(resource, nodeRes);
     numActiveNMs++;
     numActiveNMs++;

+ 33 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/NodeLabelTestBase.java

@@ -19,9 +19,11 @@
 package org.apache.hadoop.yarn.nodelabels;
 package org.apache.hadoop.yarn.nodelabels;
 
 
 import java.util.Collection;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Map;
 import java.util.Set;
 import java.util.Set;
+import java.util.Map.Entry;
 
 
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.junit.Assert;
 import org.junit.Assert;
@@ -39,6 +41,37 @@ public class NodeLabelTestBase {
     }
     }
   }
   }
 
 
+  public static void assertLabelsToNodesEquals(Map<String, Set<NodeId>> m1,
+      ImmutableMap<String, Set<NodeId>> m2) {
+    Assert.assertEquals(m1.size(), m2.size());
+    for (String k : m1.keySet()) {
+      Assert.assertTrue(m2.containsKey(k));
+      Set<NodeId> s1 = new HashSet<NodeId>(m1.get(k));
+      Set<NodeId> s2 = new HashSet<NodeId>(m2.get(k));
+      Assert.assertEquals(s1, s2);
+      Assert.assertTrue(s1.containsAll(s2));
+    }
+  }
+
+  public static ImmutableMap<String, Set<NodeId>> transposeNodeToLabels(
+      Map<NodeId, Set<String>> mapNodeToLabels) {
+    Map<String, Set<NodeId>> mapLabelsToNodes =
+        new HashMap<String, Set<NodeId>>();
+    for(Entry<NodeId, Set<String>> entry : mapNodeToLabels.entrySet()) {
+      NodeId node = entry.getKey();
+      Set<String> setLabels = entry.getValue();
+      for(String label : setLabels) {
+        Set<NodeId> setNode = mapLabelsToNodes.get(label);
+        if (setNode == null) {
+          setNode = new HashSet<NodeId>();
+        }
+        setNode.add(NodeId.newInstance(node.getHost(), node.getPort()));
+        mapLabelsToNodes.put(label, setNode);
+      }
+    }
+    return ImmutableMap.copyOf(mapLabelsToNodes);
+  }
+
   public static void assertMapContains(Map<NodeId, Set<String>> m1,
   public static void assertMapContains(Map<NodeId, Set<String>> m1,
       ImmutableMap<NodeId, Set<String>> m2) {
       ImmutableMap<NodeId, Set<String>> m2) {
     for (NodeId k : m2.keySet()) {
     for (NodeId k : m2.keySet()) {

+ 121 - 1
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestCommonNodeLabelsManager.java

@@ -18,6 +18,8 @@
 
 
 package org.apache.hadoop.yarn.nodelabels;
 package org.apache.hadoop.yarn.nodelabels;
 
 
+import static org.junit.Assert.assertTrue;
+
 import java.io.IOException;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.HashSet;
@@ -310,7 +312,6 @@ public class TestCommonNodeLabelsManager extends NodeLabelTestBase {
     
     
     // Set labels on n1:1 to P2 again to verify if add/remove works
     // Set labels on n1:1 to P2 again to verify if add/remove works
     mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p2")));
     mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p2")));
-    
     // Add p3 to n1, should makes n1:1 to be p2/p3, and n1:2 to be p1/p3
     // Add p3 to n1, should makes n1:1 to be p2/p3, and n1:2 to be p1/p3
     mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
     mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
     assertMapEquals(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n1"),
     assertMapEquals(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n1"),
@@ -397,5 +398,124 @@ public class TestCommonNodeLabelsManager extends NodeLabelTestBase {
     caught = false;
     caught = false;
     
     
     mgr.close();
     mgr.close();
+  }  
+
+  @Test(timeout = 5000)
+  public void testLabelsToNodes()
+      throws IOException {
+    mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
+    mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1"), toSet("p1", "p2")));
+    Map<String, Set<NodeId>> labelsToNodes = mgr.getLabelsToNodes();
+    assertLabelsToNodesEquals(
+        labelsToNodes,
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1")),
+        "p2",toSet(toNodeId("n1"))));
+    assertLabelsToNodesEquals(
+        labelsToNodes, transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Replace labels on n1:1 to P2
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p2"),
+        toNodeId("n1:2"), toSet("p2")));
+    labelsToNodes = mgr.getLabelsToNodes();
+    assertLabelsToNodesEquals(
+        labelsToNodes,
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1")),
+        "p2", toSet(toNodeId("n1"),toNodeId("n1:1"),toNodeId("n1:2"))));
+    assertLabelsToNodesEquals(
+        labelsToNodes, transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Replace labels on n1 to P1, both n1:1/n1 will be P1 now
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), toSet("p1")));
+    labelsToNodes = mgr.getLabelsToNodes();
+    assertLabelsToNodesEquals(
+        labelsToNodes,
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1"),toNodeId("n1:1"),toNodeId("n1:2"))));
+    assertLabelsToNodesEquals(
+        labelsToNodes, transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Set labels on n1:1 to P2 again to verify if add/remove works
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p2")));
+    // Add p3 to n1, should makes n1:1 to be p2/p3, and n1:2 to be p1/p3
+    mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
+    labelsToNodes = mgr.getLabelsToNodes();
+    assertLabelsToNodesEquals(
+        labelsToNodes,
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1"),toNodeId("n1:2")),
+        "p2", toSet(toNodeId("n1:1")),
+        "p3", toSet(toNodeId("n1"),toNodeId("n1:1"),toNodeId("n1:2"))));
+    assertLabelsToNodesEquals(
+        labelsToNodes, transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Remove P3 from n1, should makes n1:1 to be p2, and n1:2 to be p1
+    mgr.removeLabelsFromNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
+    labelsToNodes = mgr.getLabelsToNodes();
+    assertLabelsToNodesEquals(
+        labelsToNodes,
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1"),toNodeId("n1:2")),
+        "p2", toSet(toNodeId("n1:1"))));
+    assertLabelsToNodesEquals(
+        labelsToNodes, transposeNodeToLabels(mgr.getNodeLabels()));
+  }
+
+  @Test(timeout = 5000)
+  public void testLabelsToNodesForSelectedLabels()
+      throws IOException {
+    mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
+    mgr.addLabelsToNode(
+        ImmutableMap.of(
+        toNodeId("n1:1"), toSet("p1", "p2"),
+        toNodeId("n1:2"), toSet("p1", "p2")));
+    Set<String> setlabels =
+        new HashSet<String>(Arrays.asList(new String[]{"p1"}));
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of("p1", toSet(toNodeId("n1:1"),toNodeId("n1:2"))));
+
+    // Replace labels on n1:1 to P2
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
+    assertTrue(mgr.getLabelsToNodes(setlabels).isEmpty());
+    setlabels = new HashSet<String>(Arrays.asList(new String[]{"p2", "p3"}));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of(
+        "p3", toSet(toNodeId("n1"), toNodeId("n1:1"),toNodeId("n1:2"))));
+
+    mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1:3"), toSet("p1", "p2")));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of(
+        "p2", toSet(toNodeId("n1:3")),
+        "p3", toSet(toNodeId("n1"), toNodeId("n1:1"),toNodeId("n1:2"))));
+
+    mgr.removeLabelsFromNode(ImmutableMap.of(toNodeId("n1"), toSet("p3")));
+    setlabels =
+        new HashSet<String>(Arrays.asList(new String[]{"p1", "p2", "p3"}));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1:3")),
+        "p2", toSet(toNodeId("n1:3"))));
+
+    mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n2:2"), toSet("p1", "p2")));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1:3"), toNodeId("n2:2")),
+        "p2", toSet(toNodeId("n1:3"), toNodeId("n2:2"))));
+
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n2:2"), toSet("p3")));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of(
+        "p1", toSet(toNodeId("n1:3")),
+        "p2", toSet(toNodeId("n1:3")),
+        "p3", toSet(toNodeId("n2:2"))));
+    setlabels = new HashSet<String>(Arrays.asList(new String[]{"p1"}));
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(setlabels),
+        ImmutableMap.of("p1", toSet(toNodeId("n1:3"))));
   }
   }
 }
 }

+ 21 - 1
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestFileSystemNodeLabelsStore.java

@@ -116,6 +116,11 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase {
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toNodeId("n7"), toSet("p6")));
         toNodeId("n7"), toSet("p6")));
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(),
+            ImmutableMap.of(
+            "p6", toSet(toNodeId("n6"), toNodeId("n7")),
+            "p4", toSet(toNodeId("n4")),
+            "p2", toSet(toNodeId("n2"))));
 
 
     // stutdown mgr and start a new mgr
     // stutdown mgr and start a new mgr
     mgr.stop();
     mgr.stop();
@@ -130,6 +135,11 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase {
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toNodeId("n7"), toSet("p6")));
         toNodeId("n7"), toSet("p6")));
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(),
+            ImmutableMap.of(
+            "p6", toSet(toNodeId("n6"), toNodeId("n7")),
+            "p4", toSet(toNodeId("n4")),
+            "p2", toSet(toNodeId("n2"))));
     mgr.stop();
     mgr.stop();
   }
   }
 
 
@@ -169,6 +179,11 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase {
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toNodeId("n7"), toSet("p6")));
         toNodeId("n7"), toSet("p6")));
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(),
+            ImmutableMap.of(
+            "p6", toSet(toNodeId("n6"), toNodeId("n7")),
+            "p4", toSet(toNodeId("n4")),
+            "p2", toSet(toNodeId("n2"))));
     mgr.stop();
     mgr.stop();
   }
   }
   
   
@@ -218,7 +233,12 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase {
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
     assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), toSet("p6"),
         toNodeId("n7"), toSet("p6")));
         toNodeId("n7"), toSet("p6")));
-    
+    assertLabelsToNodesEquals(mgr.getLabelsToNodes(),
+        ImmutableMap.of(
+        "p6", toSet(toNodeId("n6"), toNodeId("n7")),
+        "p4", toSet(toNodeId("n4")),
+        "p2", toSet(toNodeId("n2"))));
+
     /*
     /*
      * Add label p7,p8 then shutdown
      * Add label p7,p8 then shutdown
      */
      */

+ 11 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/RMNodeLabelsManager.java

@@ -200,6 +200,17 @@ public class RMNodeLabelsManager extends CommonNodeLabelsManager {
       Node nm = getNMInNodeSet(nodeId);
       Node nm = getNMInNodeSet(nodeId);
       nm.resource = resource;
       nm.resource = resource;
       nm.running = true;
       nm.running = true;
+
+      // Add node in labelsCollection
+      Set<String> labelsForNode = getLabelsByNode(nodeId);
+      if (labelsForNode != null) {
+        for (String label : labelsForNode) {
+          NodeLabel labelInfo = labelCollections.get(label);
+          if(labelInfo != null) {
+            labelInfo.addNodeId(nodeId);
+          }
+        }
+      }
       
       
       // get the node after edition
       // get the node after edition
       Map<String, Host> after = cloneNodeMap(ImmutableSet.of(nodeId));
       Map<String, Host> after = cloneNodeMap(ImmutableSet.of(nodeId));

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

@@ -464,4 +464,31 @@ public class TestRMNodeLabelsManager extends NodeLabelTestBase {
     checkNodeLabelInfo(infos, "y", 1, 10);
     checkNodeLabelInfo(infos, "y", 1, 10);
     checkNodeLabelInfo(infos, "z", 0, 0);
     checkNodeLabelInfo(infos, "z", 0, 0);
   }
   }
+  
+  @Test(timeout = 5000)
+  public void testLabelsToNodesOnNodeActiveDeactive() throws Exception {
+    // Activate a node without assigning any labels
+    mgr.activateNode(NodeId.newInstance("n1", 1), Resource.newInstance(10, 0));
+    Assert.assertTrue(mgr.getLabelsToNodes().isEmpty());
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(), transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Add labels and replace labels on node
+    mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
+    mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), toSet("p1"),
+        toNodeId("n2"), toSet("p2"), toNodeId("n3"), toSet("p3")));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(), transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Activate a node for which host to label mapping exists
+    mgr.activateNode(NodeId.newInstance("n1", 2), Resource.newInstance(10, 0));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(), transposeNodeToLabels(mgr.getNodeLabels()));
+
+    // Deactivate a node. Label mapping should still exist.
+    mgr.deactivateNode(NodeId.newInstance("n1", 1));
+    assertLabelsToNodesEquals(
+        mgr.getLabelsToNodes(), transposeNodeToLabels(mgr.getNodeLabels()));
+  }
+
 }
 }