YARN-10501. Can't remove all node labels after add node label without

nodemanager port. Contributed by caozhiqiang.

(cherry picked from commit 4891e68c2b)
This commit is contained in:
Eric Badger 2021-02-19 23:44:28 +00:00
parent 1b5de609ec
commit 9327039fb8
2 changed files with 65 additions and 0 deletions

View File

@ -560,6 +560,50 @@ public class CommonNodeLabelsManager extends AbstractService {
addNodeToLabels(node, newLabels);
}
private void addLabelsToNodeInHost(NodeId node, Set<String> labels)
throws IOException {
Host host = nodeCollections.get(node.getHost());
if (null == host) {
throw new IOException("Cannot add labels to a host that "
+ "does not exist. Create the host before adding labels to it.");
}
Node nm = host.nms.get(node);
if (nm != null) {
Node newNm = nm.copy();
if (newNm.labels == null) {
newNm.labels =
Collections.newSetFromMap(new ConcurrentHashMap<String, Boolean>());
}
newNm.labels.addAll(labels);
host.nms.put(node, newNm);
}
}
protected void removeLabelsFromNodeInHost(NodeId node, Set<String> labels)
throws IOException {
Host host = nodeCollections.get(node.getHost());
if (null == host) {
throw new IOException("Cannot remove labels from a host that "
+ "does not exist. Create the host before adding labels to it.");
}
Node nm = host.nms.get(node);
if (nm != null) {
if (nm.labels == null) {
nm.labels = new HashSet<String>();
} else {
nm.labels.removeAll(labels);
}
}
}
private void replaceLabelsForNode(NodeId node, Set<String> oldLabels,
Set<String> newLabels) throws IOException {
if(oldLabels != null) {
removeLabelsFromNodeInHost(node, oldLabels);
}
addLabelsToNodeInHost(node, newLabels);
}
@SuppressWarnings("unchecked")
protected void internalUpdateLabelsOnNodes(
Map<NodeId, Set<String>> nodeToLabels, NodeLabelUpdateOperation op)
@ -598,10 +642,12 @@ public class CommonNodeLabelsManager extends AbstractService {
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;
@ -626,6 +672,7 @@ public class CommonNodeLabelsManager extends AbstractService {
case REPLACE:
oldLabels = getLabelsByNode(nodeId);
replaceNodeForLabels(nodeId, oldLabels, labels);
replaceLabelsForNode(nodeId, oldLabels, labels);
if (nm.labels == null) {
nm.labels = new HashSet<String>();
}

View File

@ -616,4 +616,22 @@ public class TestCommonNodeLabelsManager extends NodeLabelTestBase {
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.addToCluserNodeLabels(Arrays.asList(NodeLabel.newInstance("p2", true)));
mgr.addLabelsToNode(ImmutableMap.of(toNodeId("n1:1"), toSet("p1")));
mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), toSet("p2")));
Map<String, Set<NodeId>> labelsToNodes = mgr.getLabelsToNodes();
assertLabelsToNodesEquals(
labelsToNodes,
ImmutableMap.of(
"p2", toSet(toNodeId("n1:1"), toNodeId("n1:0"))));
mgr.replaceLabelsOnNode(ImmutableMap.of(toNodeId("n1"), new HashSet()));
Map<String, Set<NodeId>> labelsToNodes2 = mgr.getLabelsToNodes();
Assert.assertEquals(labelsToNodes2.get("p2"), null);
}
}