From 7d88953768411ef3e1187eaca039a2b0255ad091 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Tue, 2 Feb 2021 09:07:25 +0100
Subject: [PATCH] YARN-10604. Support auto queue creation without mapping
 rules. Contributed by Andras Gyori

---
 .../scheduler/capacity/CapacityScheduler.java | 18 +++++++++--
 ...CapacitySchedulerNewQueueAutoCreation.java | 30 ++++++++++++++++++-
 2 files changed, 45 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/CapacityScheduler.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/CapacityScheduler.java
index 204fa73b45..1bb74a092e 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/CapacityScheduler.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/CapacityScheduler.java
@@ -932,12 +932,26 @@ private CSQueue getOrCreateQueueFromPlacementContext(ApplicationId
       boolean isRecovery) {
 
     CSQueue queue = getQueue(queueName);
+    ApplicationPlacementContext fallbackContext = placementContext;
 
     if (queue == null) {
-      if (placementContext != null && placementContext.hasParentQueue()) {
+      // Even if placement rules are turned off, we still have the opportunity
+      // to auto create a queue.
+      if (placementContext == null) {
+        fallbackContext = CSQueueUtils.extractQueuePath(queueName);
+      }
+      if (fallbackContext.hasParentQueue()) {
         try {
-          return autoCreateLeafQueue(placementContext);
+          return autoCreateLeafQueue(fallbackContext);
         } catch (YarnException | IOException e) {
+          // A null queue is expected if the placementContext is null. In order
+          // not to disrupt the control flow, if we fail to auto create a queue,
+          // we fall back to the original logic.
+          if (placementContext == null) {
+            LOG.error("Could not auto-create leaf queue " + queueName +
+                " due to : ", e);
+            return null;
+          }
           if (isRecovery) {
             if (!getConfiguration().shouldAppFailFast(getConfig())) {
               LOG.error("Could not auto-create leaf queue " + queueName +
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/TestCapacitySchedulerNewQueueAutoCreation.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/TestCapacitySchedulerNewQueueAutoCreation.java
index 0c5375e83e..48dba82143 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/TestCapacitySchedulerNewQueueAutoCreation.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/TestCapacitySchedulerNewQueueAutoCreation.java
@@ -18,14 +18,21 @@
 
 package org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity;
 
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.resourcemanager.MockNM;
 import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
 import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.NullRMNodeLabelsManager;
 import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager;
+import org.apache.hadoop.yarn.server.resourcemanager.placement.ApplicationPlacementContext;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerDynamicEditException;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.AppAddedSchedulerEvent;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.AppAttemptAddedSchedulerEvent;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.SchedulerEvent;
+import org.apache.hadoop.yarn.server.utils.BuilderUtils;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -461,7 +468,28 @@ public void testChildlessParentQueueWhenAutoQueueCreationEnabled()
             "for auto queue creation",
         ((ParentQueue)empty).isEligibleForAutoQueueCreation());
   }
-  
+
+  @Test
+  public void testAutoQueueCreationWithDisabledMappingRules() throws Exception {
+    startScheduler();
+
+    ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
+    // Set ApplicationPlacementContext to null in the submitted application
+    // in order to imitate a submission with mapping rules turned off
+    SchedulerEvent addAppEvent = new AppAddedSchedulerEvent(appId,
+        "root.a.a1-auto.a2-auto", USER0, null);
+    ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
+        appId, 1);
+    SchedulerEvent addAttemptEvent = new AppAttemptAddedSchedulerEvent(
+        appAttemptId, false);
+    cs.handle(addAppEvent);
+    cs.handle(addAttemptEvent);
+
+    CSQueue a2Auto = cs.getQueue("root.a.a1-auto.a2-auto");
+    Assert.assertNotNull(a2Auto);
+  }
+
+  @Test
   public void testAutoCreateQueueUserLimitDisabled() throws Exception {
     startScheduler();
     createBasicQueueStructureAndValidate();