From bf570bd4acd9ebccada80a68aa1c5fbf73ca60bf Mon Sep 17 00:00:00 2001
From: 9uapaw <gyorandr@gmail.com>
Date: Thu, 28 Jul 2022 17:32:10 +0200
Subject: [PATCH] YARN-11063. Support auto queue creation template wildcards
 for arbitrary queue depths. Contributed by Bence Kosztolnik.

---
 .../capacity/AutoCreatedQueueTemplate.java       |  7 +++----
 .../capacity/TestAutoCreatedQueueTemplate.java   | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 4 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 8fa1e130d4..bbeca4b072 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
@@ -42,7 +42,6 @@ public class AutoCreatedQueueTemplate {
       AUTO_QUEUE_CREATION_V2_PREFIX + "parent-template.";
 
   private static final String WILDCARD_QUEUE = "*";
-  private static final int MAX_WILDCARD_LEVEL = 1;
 
   private final Map<String, String> templateProperties = new HashMap<>();
   private final Map<String, String> leafOnlyProperties = new HashMap<>();
@@ -86,7 +85,7 @@ public Map<String, String> getParentOnlyProperties() {
   /**
    * Sets the common template properties and parent specific template
    * properties of a child queue based on its parent template settings.
- * @param conf configuration to set
+   * @param conf configuration to set
    * @param childQueuePath child queue path used for prefixing the properties
    */
   public void setTemplateEntriesForChild(CapacitySchedulerConfiguration conf,
@@ -169,8 +168,8 @@ private void setTemplateConfigEntries(CapacitySchedulerConfiguration configurati
     // start with the most explicit format (without wildcard)
     int wildcardLevel = 0;
     // root can not be wildcarded
-    // MAX_WILDCARD_LEVEL will be configurable in the future
-    int supportedWildcardLevel = Math.min(queuePathMaxIndex, MAX_WILDCARD_LEVEL);
+    int supportedWildcardLevel = Math.min(queuePathMaxIndex,
+            configuration.getMaximumAutoCreatedQueueDepth(queuePath.getFullPath()));
     // Allow root to have template properties
     if (queuePath.isRoot()) {
       supportedWildcardLevel = 0;
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 e73001b7e5..201214ec26 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
@@ -66,6 +66,22 @@ public void testOneLevelWildcardTemplate() {
 
   }
 
+  @Test
+  public void testTwoLevelWildcardTemplate() {
+    conf.set(getTemplateKey("root.*", "capacity"), "6w");
+    conf.set(getTemplateKey("root.*.*", "capacity"), "5w");
+
+    new AutoCreatedQueueTemplate(conf, TEST_QUEUE_A)
+            .setTemplateEntriesForChild(conf, TEST_QUEUE_AB.getFullPath());
+    new AutoCreatedQueueTemplate(conf, TEST_QUEUE_AB)
+            .setTemplateEntriesForChild(conf, TEST_QUEUE_ABC.getFullPath());
+
+    Assert.assertEquals("weight is not set", 6f,
+            conf.getNonLabeledQueueWeight(TEST_QUEUE_AB.getFullPath()), 10e-6);
+    Assert.assertEquals("weight is not set", 5f,
+            conf.getNonLabeledQueueWeight(TEST_QUEUE_ABC.getFullPath()), 10e-6);
+  }
+
   @Test
   public void testIgnoredWhenRootWildcarded() {
     conf.set(getTemplateKey("*", "capacity"), "6w");