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 4fd4bd6bd72..39eb72012b3 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 @@ -559,6 +559,48 @@ private void replaceNodeForLabels(NodeId node, Set oldLabels, addNodeToLabels(node, newLabels); } + private void addLabelsToNodeInHost(NodeId node, Set labels) + throws IOException { + Host host = nodeCollections.get(node.getHost()); + if (null == host) { + throw new IOException("Should create host before add labels."); + } + Node nm = host.nms.get(node); + if (nm != null) { + Node newNm = nm.copy(); + if (newNm.labels == null) { + newNm.labels = + Collections.newSetFromMap(new ConcurrentHashMap()); + } + newNm.labels.addAll(labels); + host.nms.put(node, newNm); + } + } + + protected void removeLabelsFromNodeInHost(NodeId node, Set labels) + throws IOException { + Host host = nodeCollections.get(node.getHost()); + if (null == host) { + throw new IOException("Should create host before remove labels."); + } + Node nm = host.nms.get(node); + if (nm != null) { + if (nm.labels == null) { + nm.labels = new HashSet(); + } else { + nm.labels.removeAll(labels); + } + } + } + + private void replaceLabelsForNode(NodeId node, Set oldLabels, + Set newLabels) throws IOException { + if(oldLabels != null) { + removeLabelsFromNodeInHost(node, oldLabels); + } + addLabelsToNodeInHost(node, newLabels); + } + @SuppressWarnings("unchecked") protected void internalUpdateLabelsOnNodes( Map> nodeToLabels, NodeLabelUpdateOperation op) @@ -597,10 +639,12 @@ protected void internalUpdateLabelsOnNodes( break; case REPLACE: replaceNodeForLabels(nodeId, host.labels, labels); + replaceLabelsForNode(nodeId, host.labels, labels); host.labels.clear(); host.labels.addAll(labels); for (Node node : host.nms.values()) { replaceNodeForLabels(node.nodeId, node.labels, labels); + replaceLabelsForNode(node.nodeId, node.labels, labels); node.labels = null; } break; @@ -625,6 +669,7 @@ protected void internalUpdateLabelsOnNodes( case REPLACE: oldLabels = getLabelsByNode(nodeId); replaceNodeForLabels(nodeId, oldLabels, labels); + replaceLabelsForNode(nodeId, oldLabels, labels); if (nm.labels == null) { nm.labels = new HashSet(); } 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 e86cca4e245..472bce88ecb 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 @@ -616,4 +616,16 @@ public void testGetNodeLabelsInfo() throws IOException { toNodeId("n1"), toSet(NodeLabel.newInstance("p2", true)), toNodeId("n2"), toSet(NodeLabel.newInstance("p3", false)))); } + + @Test(timeout = 5000) + public void testRemoveNodeLabelsInfo() throws IOException { + mgr.addToCluserNodeLabels(Arrays.asList(NodeLabel.newInstance("p1", true))); + mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1"), toSet("p1"))); + + assertLabelInfoMapEquals(mgr.getNodeLabelsInfo(), ImmutableMap.of( + toNodeId("n1"), toSet(NodeLabel.newInstance("p1", true)))); + + mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), new HashSet())); + Assert.assertEquals(mgr.getNodeLabelsInfo().size(), 0); + } }