From b7ac85259c7d20c33bef9c9cb40b8aabcab70755 Mon Sep 17 00:00:00 2001
From: Rohith Sharma K S <rohithsharmaks@apache.org>
Date: Thu, 12 May 2016 15:17:49 +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     | 14 ++++++++++++++
 .../TestApplicationMasterLauncher.java             |  5 +++++
 3 files changed, 26 insertions(+)

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 d75297a4168..fa475e7712d 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
@@ -159,6 +159,13 @@ public interface ApplicationConstants {
      */
     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 b927bb42572..a160741d0b5 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.ContainerId;
 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;
@@ -191,12 +192,25 @@ public class AMLauncher implements Runnable {
         + 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 8cdb191d94d..13772dd1a68 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 @@ public class TestApplicationMasterLauncher {
     String nmHostAtContainerManager = null;
     long submitTimeAtContainerManager;
     int maxAppAttempts;
+    private String queueName;
 
     @Override
     public StartContainersResponse
@@ -120,6 +121,8 @@ public class TestApplicationMasterLauncher {
       submitTimeAtContainerManager =
           Long.parseLong(env.get(ApplicationConstants.APP_SUBMIT_TIME_ENV));
       maxAppAttempts = YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS;
+      queueName = env.get(ApplicationConstants.Environment
+              .YARN_RESOURCEMANAGER_APPLICATION_QUEUE.key());
       return StartContainersResponse.newInstance(
         new HashMap<String, ByteBuffer>(), new ArrayList<ContainerId>(),
         new HashMap<ContainerId, SerializedException>());
@@ -182,6 +185,8 @@ public class TestApplicationMasterLauncher {
       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);