YARN-10835. Pass user defined variables of yarn.nodemanager.env-whitelist along jobs

hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/main/java/org/apache/hadoop/mapreduce/v2/app/job/impl/TaskAttemptImpl.java

@@ -1137,6 +1137,17 @@ static ContainerLaunchContext createContainerLaunchContext(
     }
     MapReduceChildJVM.setVMEnv(myEnv, remoteTask);
 
+    if (conf.getBoolean(YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED,
+        YarnConfiguration.DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED)) {
+      String[] whitelistVars = conf.get(YarnConfiguration.NM_ENV_WHITELIST,
+          YarnConfiguration.DEFAULT_NM_ENV_WHITELIST).split(",");
+      for (String var : whitelistVars) {
+        if (System.getenv(var) != null && !System.getenv(var).isEmpty()) {
+          myEnv.put(var, System.getenv(var));
+        }
+      }
+    }
+
     // Set up the launch command
     List<String> commands = MapReduceChildJVM.getVMCommand(
         taskAttemptListener.getAddress(), remoteTask, jvmID);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java

@@ -1265,7 +1265,10 @@ public static boolean isAclEnabled(Configuration conf) {
       ApplicationConstants.Environment.HADOOP_CONF_DIR.key(),
       ApplicationConstants.Environment.CLASSPATH_PREPEND_DISTCACHE.key(),
       ApplicationConstants.Environment.HADOOP_YARN_HOME.key()));
-
+  public static final String NM_ENV_WHITELIST_EXPORT_ENABLED = NM_ENV_WHITELIST
+      + "export.enabled";
+  public static final boolean DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED = false;
+
   /** address of node manager IPC.*/
   public static final String NM_ADDRESS = NM_PREFIX + "address";
   public static final int DEFAULT_NM_PORT = 0;

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java

@@ -310,6 +310,20 @@ public YarnClientApplication createApplication()
       throw new ApplicationIdNotProvidedException(
           "ApplicationId is not provided in ApplicationSubmissionContext");
     }
+
+    if (isNMEnvWhitelistExportEnabled()) {
+      String[] whitelistVars = getConfig().get(
+          YarnConfiguration.NM_ENV_WHITELIST,
+          YarnConfiguration.DEFAULT_NM_ENV_WHITELIST).split(",");
+      for(String var : whitelistVars) {
+        String value = getNMEnvWhitelistValue(var);
+        if (value != null) {
+          appContext.getAMContainerSpec().getEnvironment()
+              .put(var, value);
+        }
+      }
+    }
+
     SubmitApplicationRequest request =
         Records.newRecord(SubmitApplicationRequest.class);
     request.setApplicationSubmissionContext(appContext);
@@ -524,6 +538,12 @@ protected boolean isSecurityEnabled() {
     return UserGroupInformation.isSecurityEnabled();
   }
 
+  protected boolean isNMEnvWhitelistExportEnabled() {
+    return getConfig().getBoolean(
+        YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED,
+        YarnConfiguration.DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED);
+  }
+
   @Override
   public void failApplicationAttempt(ApplicationAttemptId attemptId)
       throws YarnException, IOException {
@@ -1191,4 +1211,14 @@ public void shellToContainer(ContainerId containerId,
       LOG.error("Fail to shell to container: " + t.getMessage());
     }
   }
+
+  private String getNMEnvWhitelistValue(String key) {
+    if (System.getenv(key) != null && !System.getenv(key).isEmpty()) {
+      return System.getenv(key);
+    }
+    if (System.getProperty(key) != null && !System.getProperty(key).isEmpty()) {
+      return System.getProperty(key);
+    }
+    return null;
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClientImpl.java

@@ -59,6 +59,8 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
@@ -381,4 +383,63 @@ public void testParseTimelineDelegationTokenRenewer() {
       client.stop();
     }
   }
+
+  @Test
+  public void testNMEnvWhitelistExportEnabled() {
+    Configuration conf = getConf();
+    conf.set(YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED, "true");
+    conf.set(YarnConfiguration.NM_ENV_WHITELIST, "TEST_KEY_1");
+
+    // Set property for TEST_KEY_1
+    System.setProperty("TEST_KEY_1", "TEST_VALUE_1");
+    System.setProperty("TEST_KEY_2", "TEST_VALUE_2");
+
+    // Prepare a Mock YarnClientImpl
+    YarnClientImpl client = spy(new YarnClientImpl() {
+
+      @Override
+      protected void serviceStart() {
+        rmClient = mock(ApplicationClientProtocol.class);
+      }
+
+      @Override
+      protected void serviceStop() {
+      }
+
+      @Override
+      public ApplicationReport getApplicationReport(ApplicationId appId) {
+        ApplicationReport report = mock(ApplicationReport.class);
+        when(report.getYarnApplicationState())
+            .thenReturn(YarnApplicationState.RUNNING);
+        return report;
+      }
+    });
+    conf.set(
+        YarnConfiguration.RM_ADDRESS, "localhost:8188");
+
+    // Prepare a ApplicationSubmissionContext and submit the app
+    ApplicationSubmissionContext context =
+        mock(ApplicationSubmissionContext.class);
+    ApplicationId applicationId = ApplicationId.newInstance(0, 1);
+    when(context.getApplicationId()).thenReturn(applicationId);
+
+    Map<String, String> env = new HashMap<>();
+    ContainerLaunchContext clc = ContainerLaunchContext.newInstance(
+        null, env, null, null, null, null);
+    when(context.getAMContainerSpec()).thenReturn(clc);
+
+    try {
+      client.init(conf);
+      client.start();
+      client.submitApplication(context);
+      Assert.assertEquals("TEST_VALUE_1",
+          context.getAMContainerSpec().getEnvironment().get("TEST_KEY_1"));
+      Assert.assertNotEquals("TEST_VALUE_2",
+          context.getAMContainerSpec().getEnvironment().get("TEST_KEY_2"));
+    } catch (YarnException | IOException e) {
+      e.printStackTrace();
+    } finally {
+      client.stop();
+    }
+  }
 }