From 9327039fb8d1cb27beab49a8da14b4136265257c Mon Sep 17 00:00:00 2001
From: Eric Badger <ebadger@verizonmedia.com>
Date: Fri, 19 Feb 2021 23:44:28 +0000
Subject: [PATCH] YARN-10501. Can't remove all node labels after add node label
 without nodemanager port. Contributed by caozhiqiang.

(cherry picked from commit 4891e68c2b2bdf36589da1174d4d8ee29acc57f6)
---
 .../nodelabels/CommonNodeLabelsManager.java   | 47 +++++++++++++++++++
 .../TestCommonNodeLabelsManager.java          | 18 +++++++
 2 files changed, 65 insertions(+)

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 66e945fc2ff..88ffb5d16f4 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
@@ -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>();
             }
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 54e331b71b9..6b4bb6d786c 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,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);
+  }
 }