From f102c4f6c8b5a38f73e41282215e9bc38240eed8 Mon Sep 17 00:00:00 2001
From: Rohith Sharma K S <rohithsharmaks@apache.org>
Date: Mon, 16 May 2016 15:25:15 +0530
Subject: [PATCH] YARN-5068. Expose scheduler queue to application master.
 (Harish Jaiprakash via rohithsharmaks)

---
 .../hadoop/yarn/api/ApplicationConstants.java  |  7 +++++++
 .../resourcemanager/amlauncher/AMLauncher.java | 18 ++++++++++++++++--
 .../TestApplicationMasterLauncher.java         |  5 +++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/ApplicationConstants.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/ApplicationConstants.java
index b2d765ae308..ad526d62afc 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/ApplicationConstants.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/ApplicationConstants.java
@@ -170,6 +170,13 @@ public enum Environment {
      */
     LD_LIBRARY_PATH("LD_LIBRARY_PATH"),
     
+    /**
+     * $YARN_RESOURCEMANAGER_APPLICATION_QUEUE
+     * The queue into which the app was submitted/launched.
+     */
+    YARN_RESOURCEMANAGER_APPLICATION_QUEUE(
+        "YARN_RESOURCEMANAGER_APPLICATION_QUEUE"),
+
     /**
      * $HADOOP_CONF_DIR
      * Final, non-modifiable.
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/amlauncher/AMLauncher.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/amlauncher/AMLauncher.java
index d2fa587197d..4c840e74a47 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/amlauncher/AMLauncher.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/amlauncher/AMLauncher.java
@@ -51,6 +51,7 @@
 import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.client.NMProxy;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.event.EventHandler;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.ipc.YarnRPC;
@@ -190,13 +191,26 @@ private ContainerLaunchContext createAMContainerLaunchContext(
         + " : "
         + StringUtils.arrayToString(container.getCommands().toArray(
             new String[0])));
-    
+
+    // Populate the current queue name in the environment variable.
+    setupQueueNameEnv(container, applicationMasterContext);
+
     // Finalize the container
     setupTokens(container, containerID);
-    
+
     return container;
   }
 
+  private void setupQueueNameEnv(ContainerLaunchContext container,
+      ApplicationSubmissionContext applicationMasterContext) {
+    String queueName = applicationMasterContext.getQueue();
+    if (queueName == null) {
+      queueName = YarnConfiguration.DEFAULT_QUEUE_NAME;
+    }
+    container.getEnvironment().put(ApplicationConstants.Environment
+            .YARN_RESOURCEMANAGER_APPLICATION_QUEUE.key(), queueName);
+  }
+
   @Private
   @VisibleForTesting
   protected void setupTokens(
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestApplicationMasterLauncher.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestApplicationMasterLauncher.java
index 01359b5f8e9..6ade92d3c30 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestApplicationMasterLauncher.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestApplicationMasterLauncher.java
@@ -92,6 +92,7 @@ private static final class MyContainerManagerImpl implements
     String nmHostAtContainerManager = null;
     long submitTimeAtContainerManager;
     int maxAppAttempts;
+    private String queueName;
 
     @Override
     public StartContainersResponse
@@ -121,6 +122,8 @@ private static final class MyContainerManagerImpl implements
           Long.parseLong(env.get(ApplicationConstants.APP_SUBMIT_TIME_ENV));
       maxAppAttempts =
           Integer.parseInt(env.get(ApplicationConstants.MAX_APP_ATTEMPTS_ENV));
+      queueName = env.get(ApplicationConstants.Environment
+              .YARN_RESOURCEMANAGER_APPLICATION_QUEUE.key());
       return StartContainersResponse.newInstance(
         new HashMap<String, ByteBuffer>(), new ArrayList<ContainerId>(),
         new HashMap<ContainerId, SerializedException>());
@@ -183,6 +186,8 @@ public void testAMLaunchAndCleanup() throws Exception {
       containerManager.nmHostAtContainerManager);
     Assert.assertEquals(YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS,
         containerManager.maxAppAttempts);
+    Assert.assertEquals(YarnConfiguration.DEFAULT_QUEUE_NAME,
+        containerManager.queueName);
 
     MockAM am = new MockAM(rm.getRMContext(), rm
         .getApplicationMasterService(), appAttemptId);