From 205e15c4a322cbf9674844cad4c806e9b9ebe48b Mon Sep 17 00:00:00 2001 From: Wangda Tan Date: Tue, 3 Feb 2015 12:51:21 -0800 Subject: [PATCH] YARN-3075. NodeLabelsManager implementation to retrieve label to node mapping (Varun Saxena via wangda) (cherry picked from commit 5bd984691b3e3df5d8aac910295138996e7a5839) --- hadoop-yarn-project/CHANGES.txt | 3 + .../nodelabels/CommonNodeLabelsManager.java | 98 ++++++++++++-- .../hadoop/yarn/nodelabels/NodeLabel.java | 17 +++ .../yarn/nodelabels/NodeLabelTestBase.java | 33 +++++ .../TestCommonNodeLabelsManager.java | 122 +++++++++++++++++- .../TestFileSystemNodeLabelsStore.java | 22 +++- .../nodelabels/RMNodeLabelsManager.java | 11 ++ .../nodelabels/TestRMNodeLabelsManager.java | 27 ++++ 8 files changed, 323 insertions(+), 10 deletions(-) diff --git a/hadoop-yarn-project/CHANGES.txt b/hadoop-yarn-project/CHANGES.txt index 859061943f3..de77bbb8e48 100644 --- a/hadoop-yarn-project/CHANGES.txt +++ b/hadoop-yarn-project/CHANGES.txt @@ -200,6 +200,9 @@ Release 2.7.0 - UNRELEASED YARN-3098. Created common QueueCapacities class in Capacity Scheduler to 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 BUG FIXES diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java index daa1529414b..f8e7a209af8 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java +++ b/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 labels; public Resource resource; public boolean running; + public NodeId nodeId; - protected Node() { + protected Node(NodeId nodeid) { labels = null; resource = Resource.newInstance(0, 0); running = false; + nodeId = nodeid; } public Node copy() { - Node c = new Node(); + Node c = new Node(nodeId); if (labels != null) { c.labels = Collections.newSetFromMap(new ConcurrentHashMap()); @@ -465,7 +467,27 @@ public class CommonNodeLabelsManager extends AbstractService { } } } - + + private void addNodeToLabels(NodeId node, Set labels) { + for(String l : labels) { + labelCollections.get(l).addNodeId(node); + } + } + + private void removeNodeFromLabels(NodeId node, Set labels) { + for(String l : labels) { + labelCollections.get(l).removeNodeId(node); + } + } + + private void replaceNodeForLabels(NodeId node, Set oldLabels, + Set newLabels) { + if(oldLabels != null) { + removeNodeFromLabels(node, oldLabels); + } + addNodeToLabels(node, newLabels); + } + @SuppressWarnings("unchecked") protected void internalUpdateLabelsOnNodes( Map> nodeToLabels, NodeLabelUpdateOperation op) @@ -473,6 +495,7 @@ public class CommonNodeLabelsManager extends AbstractService { // do update labels from nodes Map> newNMToLabels = new HashMap>(); + Set oldLabels; for (Entry> entry : nodeToLabels.entrySet()) { NodeId nodeId = entry.getKey(); Set labels = entry.getValue(); @@ -481,26 +504,32 @@ public class CommonNodeLabelsManager extends AbstractService { if (nodeId.getPort() == WILDCARD_PORT) { Host host = nodeCollections.get(nodeId.getHost()); switch (op) { - case REMOVE: + case REMOVE: + removeNodeFromLabels(nodeId, labels); host.labels.removeAll(labels); for (Node node : host.nms.values()) { if (node.labels != null) { node.labels.removeAll(labels); } + removeNodeFromLabels(node.nodeId, labels); } break; case ADD: + addNodeToLabels(nodeId, labels); host.labels.addAll(labels); for (Node node : host.nms.values()) { if (node.labels != null) { node.labels.addAll(labels); } + addNodeToLabels(node.nodeId, labels); } break; case REPLACE: + replaceNodeForLabels(nodeId, host.labels, labels); host.labels.clear(); host.labels.addAll(labels); for (Node node : host.nms.values()) { + replaceNodeForLabels(node.nodeId, node.labels, labels); node.labels = null; } break; @@ -514,14 +543,20 @@ public class CommonNodeLabelsManager extends AbstractService { // Add and replace createNodeIfNonExisted(nodeId); Node nm = getNMInNodeSet(nodeId); - if (nm.labels == null) { - nm.labels = new HashSet(); - } switch (op) { case ADD: + addNodeToLabels(nodeId, labels); + if (nm.labels == null) { + nm.labels = new HashSet(); + } nm.labels.addAll(labels); break; case REPLACE: + oldLabels = getLabelsByNode(nodeId); + replaceNodeForLabels(nodeId, oldLabels, labels); + if (nm.labels == null) { + nm.labels = new HashSet(); + } nm.labels.clear(); nm.labels.addAll(labels); break; @@ -531,6 +566,7 @@ public class CommonNodeLabelsManager extends AbstractService { newNMToLabels.put(nodeId, nm.labels); } else { // remove + removeNodeFromLabels(nodeId, labels); Node nm = getNMInNodeSet(nodeId); if (nm.labels != null) { 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> 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> getLabelsToNodes(Set labels) { + try { + readLock.lock(); + Map> labelsToNodes = + new HashMap>(); + for (String label : labels) { + if(label.equals(NO_LABEL)) { + continue; + } + NodeLabel nodeLabelInfo = labelCollections.get(label); + if(nodeLabelInfo != null) { + Set 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 * @@ -741,7 +823,7 @@ public class CommonNodeLabelsManager extends AbstractService { } Node nm = host.nms.get(nodeId); if (null == nm) { - host.nms.put(nodeId, new Node()); + host.nms.put(nodeId, new Node(nodeId)); } } diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeLabel.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeLabel.java index 766864887d0..1765a65b26c 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeLabel.java +++ b/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; +import java.util.HashSet; +import java.util.Set; 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.util.resource.Resources; @@ -26,6 +29,7 @@ public class NodeLabel implements Comparable { private Resource resource; private int numActiveNMs; private String labelName; + private Set nodeIds; public NodeLabel(String labelName) { this(labelName, Resource.newInstance(0, 0), 0); @@ -35,8 +39,21 @@ public class NodeLabel implements Comparable { this.labelName = labelName; this.resource = res; this.numActiveNMs = activeNMs; + this.nodeIds = new HashSet(); + } + + public void addNodeId(NodeId node) { + nodeIds.add(node); + } + + public void removeNodeId(NodeId node) { + nodeIds.remove(node); } + public Set getAssociatedNodeIds() { + return new HashSet(nodeIds); + } + public void addNode(Resource nodeRes) { Resources.addTo(resource, nodeRes); numActiveNMs++; diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/NodeLabelTestBase.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/NodeLabelTestBase.java index ff0e1019b74..3b9825d1e03 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/NodeLabelTestBase.java +++ b/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; import java.util.Collection; +import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.Map.Entry; import org.apache.hadoop.yarn.api.records.NodeId; import org.junit.Assert; @@ -39,6 +41,37 @@ public class NodeLabelTestBase { } } + public static void assertLabelsToNodesEquals(Map> m1, + ImmutableMap> m2) { + Assert.assertEquals(m1.size(), m2.size()); + for (String k : m1.keySet()) { + Assert.assertTrue(m2.containsKey(k)); + Set s1 = new HashSet(m1.get(k)); + Set s2 = new HashSet(m2.get(k)); + Assert.assertEquals(s1, s2); + Assert.assertTrue(s1.containsAll(s2)); + } + } + + public static ImmutableMap> transposeNodeToLabels( + Map> mapNodeToLabels) { + Map> mapLabelsToNodes = + new HashMap>(); + for(Entry> entry : mapNodeToLabels.entrySet()) { + NodeId node = entry.getKey(); + Set setLabels = entry.getValue(); + for(String label : setLabels) { + Set setNode = mapLabelsToNodes.get(label); + if (setNode == null) { + setNode = new HashSet(); + } + setNode.add(NodeId.newInstance(node.getHost(), node.getPort())); + mapLabelsToNodes.put(label, setNode); + } + } + return ImmutableMap.copyOf(mapLabelsToNodes); + } + public static void assertMapContains(Map> m1, ImmutableMap> m2) { for (NodeId k : m2.keySet()) { diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestCommonNodeLabelsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestCommonNodeLabelsManager.java index 0ab11154918..11a2621e25e 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestCommonNodeLabelsManager.java +++ b/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; +import static org.junit.Assert.assertTrue; + import java.io.IOException; import java.util.Arrays; 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 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"))); assertMapEquals(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n1"), @@ -397,5 +398,124 @@ public class TestCommonNodeLabelsManager extends NodeLabelTestBase { caught = false; 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> 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 setlabels = + new HashSet(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(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(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(Arrays.asList(new String[]{"p1"})); + assertLabelsToNodesEquals(mgr.getLabelsToNodes(setlabels), + ImmutableMap.of("p1", toSet(toNodeId("n1:3")))); } } diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestFileSystemNodeLabelsStore.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestFileSystemNodeLabelsStore.java index 55e0e17fa7c..5cc026a74cc 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/nodelabels/TestFileSystemNodeLabelsStore.java +++ b/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"), toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), 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 mgr.stop(); @@ -130,6 +135,11 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase { assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"), toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), 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(); } @@ -169,6 +179,11 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase { assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"), toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), 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(); } @@ -218,7 +233,12 @@ public class TestFileSystemNodeLabelsStore extends NodeLabelTestBase { assertMapContains(mgr.getNodeLabels(), ImmutableMap.of(toNodeId("n2"), toSet("p2"), toNodeId("n4"), toSet("p4"), toNodeId("n6"), 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 */ diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/RMNodeLabelsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/RMNodeLabelsManager.java index 18478e36152..1555291cf19 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/RMNodeLabelsManager.java +++ b/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); nm.resource = resource; nm.running = true; + + // Add node in labelsCollection + Set 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 Map after = cloneNodeMap(ImmutableSet.of(nodeId)); diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/TestRMNodeLabelsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/TestRMNodeLabelsManager.java index 4b8fe766766..350586694a1 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/TestRMNodeLabelsManager.java +++ b/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, "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())); + } + }