From 126079612c8b567041ec3ea466e787a932bdf635 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Tue, 7 Dec 2021 15:55:16 +0100
Subject: [PATCH] YARN-11016. Queue weight is incorrectly reset to zero.
 Contributed by Andras Gyori

---
 .../scheduler/capacity/QueueCapacities.java   |  2 +-
 .../TestCapacitySchedulerWeightMode.java      | 46 +++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/QueueCapacities.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/QueueCapacities.java
index 2d829663eb4..0edf83f7269 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/QueueCapacities.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/QueueCapacities.java
@@ -312,7 +312,7 @@ public class QueueCapacities {
         _set(label, CapacityType.MAX_CAP, 0);
         _set(label, CapacityType.ABS_CAP, 0);
         _set(label, CapacityType.ABS_MAX_CAP, 0);
-        _set(label, CapacityType.WEIGHT, 0);
+        _set(label, CapacityType.WEIGHT, -1);
       }
     } finally {
       writeLock.unlock();
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCapacitySchedulerWeightMode.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCapacitySchedulerWeightMode.java
index a5d80dc27e5..115813a0171 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCapacitySchedulerWeightMode.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCapacitySchedulerWeightMode.java
@@ -26,6 +26,7 @@ import org.apache.hadoop.util.Sets;
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.NodeId;
+import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.server.resourcemanager.MockAM;
 import org.apache.hadoop.yarn.server.resourcemanager.MockNM;
@@ -38,6 +39,7 @@ import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsMana
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
 import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer;
 import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerState;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceLimits;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerAppReport;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.YarnScheduler;
@@ -74,6 +76,28 @@ public class TestCapacitySchedulerWeightMode {
     return set;
   }
 
+  public static CapacitySchedulerConfiguration getConfigWithInheritedAccessibleNodeLabel(
+      Configuration config) {
+    CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration(
+        config);
+
+    // Define top-level queues
+    conf.setQueues(CapacitySchedulerConfiguration.ROOT,
+        new String[] { "a"});
+
+    conf.setCapacityByLabel(A, RMNodeLabelsManager.NO_LABEL, 100f);
+    conf.setCapacityByLabel(A, "newLabel", 100f);
+    conf.setAccessibleNodeLabels(A, toSet("newLabel"));
+    conf.setAllowZeroCapacitySum(A, true);
+
+    // Define 2nd-level queues
+    conf.setQueues(A, new String[] { "a1" });
+    conf.setCapacityByLabel(A1, RMNodeLabelsManager.NO_LABEL, 100f);
+
+    return conf;
+  }
+
+
   /*
    * Queue structure:
    *                      root (*)
@@ -368,6 +392,28 @@ public class TestCapacitySchedulerWeightMode {
     rm.close();
   }
 
+  /**
+   * Tests whether weight is correctly reset to -1. See YARN-11016 for further details.
+   * @throws IOException if reinitialization fails
+   */
+  @Test()
+  public void testAccessibleNodeLabelsInheritanceNoWeightMode() throws IOException {
+    CapacitySchedulerConfiguration newConf = getConfigWithInheritedAccessibleNodeLabel(conf);
+
+    MockRM rm = new MockRM(newConf);
+    CapacityScheduler cs =
+        (CapacityScheduler) rm.getRMContext().getScheduler();
+
+    Resource clusterResource = Resource.newInstance(1024, 2);
+    cs.getRootQueue().updateClusterResource(clusterResource, new ResourceLimits(clusterResource));
+
+    try {
+      cs.reinitialize(newConf, rm.getRMContext());
+    } catch (Exception e) {
+      Assert.fail("Reinitialization failed with " + e);
+    }
+  }
+
   @Test
   public void testQueueInfoWeight() throws Exception {
     MockRM rm = new MockRM(conf);