From 2541efa496ba0e7e096ee5ec3c08d64b62036402 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Tue, 25 May 2021 13:55:59 +0200
Subject: [PATCH] YARN-10783. Allow definition of auto queue template
 properties in root. Contributed by Andras Gyori

---
 .../capacity/AutoCreatedQueueTemplate.java         | 14 +++++++++++---
 .../capacity/TestAutoCreatedQueueTemplate.java     | 11 +++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

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/AutoCreatedQueueTemplate.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/AutoCreatedQueueTemplate.java
index 6c516c04770..203ec4dcd6f 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/AutoCreatedQueueTemplate.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/AutoCreatedQueueTemplate.java
@@ -29,6 +29,7 @@
 import java.util.Set;
 
 import static org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerConfiguration.AUTO_QUEUE_CREATION_V2_PREFIX;
+import static org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerConfiguration.ROOT;
 
 /**
  * A handler for storing and setting auto created queue template settings.
@@ -68,6 +69,10 @@ public Map<String, String> getTemplateProperties() {
    */
   public void setTemplateEntriesForChild(Configuration conf,
                                          String childQueuePath) {
+    if (childQueuePath.equals(ROOT)) {
+      return;
+    }
+
     // Get all properties that are explicitly set
     Set<String> alreadySetProps = conf.getPropsWithPrefix(
         CapacitySchedulerConfiguration.getQueuePrefix(childQueuePath)).keySet();
@@ -94,8 +99,8 @@ private void setTemplateConfigEntries(Configuration configuration,
     List<String> queuePathParts = new ArrayList<>(Arrays.asList(
         queuePath.split("\\.")));
 
-    if (queuePathParts.size() <= 1) {
-      // This is either root or an empty queue name
+    if (queuePathParts.size() <= 1 && !queuePath.equals(ROOT)) {
+      // This is an invalid queue path
       return;
     }
     int queuePathMaxIndex = queuePathParts.size() - 1;
@@ -106,7 +111,10 @@ private void setTemplateConfigEntries(Configuration configuration,
     // MAX_WILDCARD_LEVEL will be configurable in the future
     int supportedWildcardLevel = Math.min(queuePathMaxIndex - 1,
         MAX_WILDCARD_LEVEL);
-
+    // Allow root to have template properties
+    if (queuePath.equals(ROOT)) {
+      supportedWildcardLevel = 0;
+    }
 
     // Collect all template entries
     while (wildcardLevel <= supportedWildcardLevel) {
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/TestAutoCreatedQueueTemplate.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/TestAutoCreatedQueueTemplate.java
index 1c021f6efb0..2763af026a6 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/TestAutoCreatedQueueTemplate.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/TestAutoCreatedQueueTemplate.java
@@ -109,6 +109,17 @@ public void testTemplatePrecedence() {
         newConf.getNonLabeledQueueWeight(TEST_QUEUE_ABC), 10e-6);
   }
 
+  @Test
+  public void testRootTemplate() {
+    conf.set(getTemplateKey("root", "capacity"), "2w");
+
+    AutoCreatedQueueTemplate template =
+        new AutoCreatedQueueTemplate(conf, ROOT);
+    template.setTemplateEntriesForChild(conf, TEST_QUEUE_A);
+    Assert.assertEquals("root property is not set", 2f,
+        conf.getNonLabeledQueueWeight(TEST_QUEUE_A), 10e-6);
+  }
+
   private String getTemplateKey(String queuePath, String entryKey) {
     return CapacitySchedulerConfiguration.getQueuePrefix(queuePath)
         + AutoCreatedQueueTemplate.AUTO_QUEUE_TEMPLATE_PREFIX + entryKey;