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

nodemanager port. Contributed by caozhiqiang.
This commit is contained in:
Eric Badger 2021-02-19 22:29:10 +00:00
parent d28b6f90c8
commit 4891e68c2b
2 changed files with 65 additions and 0 deletions

View File

@ -559,6 +559,50 @@ private void replaceNodeForLabels(NodeId node, Set<String> oldLabels,
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)
@ -597,10 +641,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 +671,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<String>();
}

View File

@ -616,4 +616,22 @@ 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.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);
}
}