From 6857a05d6ac566a60336c0a28951f09ecda39f24 Mon Sep 17 00:00:00 2001
From: Eric Badger <ebadger@verizonmedia.com>
Date: Fri, 23 Apr 2021 22:02:04 +0000
Subject: [PATCH] YARN-10479. Can't remove all node labels after add node label
 without nodemanager port, broken by YARN-10647. Contributed by D M Murali
 Krishna Reddy

---
 .../yarn/nodelabels/CommonNodeLabelsManager.java  | 15 ++++++++++++---
 .../nodelabels/RMNodeLabelsManager.java           |  2 +-
 2 files changed, 13 insertions(+), 4 deletions(-)

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 c19faaaad3..7dadaeb379 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
@@ -95,6 +95,8 @@ public class CommonNodeLabelsManager extends AbstractService {
       new ConcurrentHashMap<String, RMNodeLabel>();
   protected ConcurrentMap<String, Host> nodeCollections =
       new ConcurrentHashMap<String, Host>();
+  private ConcurrentMap<NodeId, Boolean> isNodeLabelFromHost =
+      new ConcurrentHashMap<NodeId, Boolean>();
 
   protected RMNodeLabel noNodeLabel;
 
@@ -603,6 +605,11 @@ private void replaceLabelsForNode(NodeId node, Set<String> oldLabels,
     addLabelsToNodeInHost(node, newLabels);
   }
 
+  protected boolean isNodeLabelExplicit(NodeId nodeId) {
+    return !isNodeLabelFromHost.containsKey(nodeId) ||
+        isNodeLabelFromHost.get(nodeId);
+  }
+
   @SuppressWarnings("unchecked")
   protected void internalUpdateLabelsOnNodes(
       Map<NodeId, Set<String>> nodeToLabels, NodeLabelUpdateOperation op)
@@ -637,6 +644,7 @@ protected void internalUpdateLabelsOnNodes(
               node.labels.addAll(labels);
             }
             addNodeToLabels(node.nodeId, labels);
+            isNodeLabelFromHost.put(node.nodeId, true);
           }
           break;
         case REPLACE:
@@ -646,10 +654,9 @@ protected void internalUpdateLabelsOnNodes(
           host.labels.addAll(labels);
           for (Node node : host.nms.values()) {
             replaceNodeForLabels(node.nodeId, node.labels, labels);
-            if (node.labels != null) {
-              replaceLabelsForNode(node.nodeId, node.labels, labels);
-            }
+            replaceLabelsForNode(node.nodeId, node.labels, labels);
             node.labels = null;
+            isNodeLabelFromHost.put(node.nodeId, true);
           }
           break;
         default:
@@ -669,6 +676,7 @@ protected void internalUpdateLabelsOnNodes(
               nm.labels = new HashSet<String>();
             }
             nm.labels.addAll(labels);
+            isNodeLabelFromHost.put(nm.nodeId, false);
             break;
           case REPLACE:
             oldLabels = getLabelsByNode(nodeId);
@@ -679,6 +687,7 @@ protected void internalUpdateLabelsOnNodes(
             }
             nm.labels.clear();
             nm.labels.addAll(labels);
+            isNodeLabelFromHost.put(nm.nodeId, false);
             break;
           default:
             break;
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 1a08e03e69..57f5803f9b 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
@@ -275,7 +275,7 @@ public void deactivateNode(NodeId nodeId) {
       Map<String, Host> before = cloneNodeMap(ImmutableSet.of(nodeId));
       Node nm = getNMInNodeSet(nodeId);
       if (null != nm) {
-        if (null == nm.labels) {
+        if (isNodeLabelExplicit(nm.nodeId)) {
           // When node deactivated, remove the nm from node collection if no
           // labels explicitly set for this particular nm