YARN-8071. Add ability to specify nodemanager environment variables i…

…ndividually. Contributed by Jim Brennan

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/util/Apps.java

@@ -23,10 +23,11 @@
 import static org.apache.hadoop.yarn.util.StringHelper.sjoin;
 
 import java.io.File;
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -191,20 +192,36 @@ public static void setEnvFromInputProperty(Map<String, String> env,
   /**
    *
    * @param envString String containing env variable definitions
-   * @param classPathSeparator String that separates the definitions
-   * @return ArrayList of environment variable names
+   * @return Set of environment variable names
    */
-  public static ArrayList<String> getEnvVarsFromInputString(String envString,
-      String classPathSeparator) {
-    ArrayList<String> envList = new ArrayList<>();
+  private static Set<String> getEnvVarsFromInputString(String envString) {
+    Set<String> envSet = new HashSet<>();
     if (envString != null && envString.length() > 0) {
       Matcher varValMatcher = VARVAL_SPLITTER.matcher(envString);
       while (varValMatcher.find()) {
         String envVar = varValMatcher.group(1);
-        envList.add(envVar);
+        envSet.add(envVar);
       }
     }
-    return envList;
+    return envSet;
+  }
+
+  /**
+   * Return the list of environment variable names specified in the
+   * given property or default string and those specified individually
+   * with the propname.VARNAME syntax (e.g., mapreduce.map.env.VARNAME=value).
+   * @param propName the name of the property
+   * @param defaultPropValue the default value for propName
+   * @param conf configuration containing properties
+   * @return Set of environment variable names
+   */
+  public static Set<String> getEnvVarsFromInputProperty(
+      String propName, String defaultPropValue, Configuration conf) {
+    String envString = conf.get(propName, defaultPropValue);
+    Set<String> varSet = getEnvVarsFromInputString(envString);
+    Map<String, String> propMap = conf.getPropsWithPrefix(propName + ".");
+    varSet.addAll(propMap.keySet());
+    return varSet;
   }
 
   /**

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/resources/yarn-default.xml

@@ -1059,7 +1059,16 @@
   </property>
 
   <property>
-    <description>Environment variables that should be forwarded from the NodeManager's environment to the container's.</description>
+    <description>
+      Environment variables that should be forwarded from the NodeManager's
+      environment to the container's, specified as a comma separated list of
+      VARNAME=value pairs.
+
+      To define environment variables individually, you can specify
+      multiple properties of the form yarn.nodemanager.admin-env.VARNAME,
+      where VARNAME is the name of the environment variable. This is the only
+      way to add a variable when its value contains commas.
+    </description>
     <name>yarn.nodemanager.admin-env</name>
     <value>MALLOC_ARENA_MAX=$MALLOC_ARENA_MAX</value>
   </property>

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/ContainerLaunch.java

@@ -1735,13 +1735,14 @@ public void sanitizeEnv(Map<String, String> environment, Path pwd,
       addToEnvMap(environment, nmVars, "JVM_PID", "$$");
     }
 
-    // variables here will be forced in, even if the container has specified them.
-    String nmAdminUserEnv = conf.get(
-        YarnConfiguration.NM_ADMIN_USER_ENV,
-        YarnConfiguration.DEFAULT_NM_ADMIN_USER_ENV);
-    Apps.setEnvFromInputString(environment, nmAdminUserEnv, File.pathSeparator);
-    nmVars.addAll(Apps.getEnvVarsFromInputString(nmAdminUserEnv,
-        File.pathSeparator));
+    // variables here will be forced in, even if the container has
+    // specified them.
+    String defEnvStr = conf.get(YarnConfiguration.DEFAULT_NM_ADMIN_USER_ENV);
+    Apps.setEnvFromInputProperty(environment,
+        YarnConfiguration.NM_ADMIN_USER_ENV, defEnvStr, conf,
+        File.pathSeparator);
+    nmVars.addAll(Apps.getEnvVarsFromInputProperty(
+        YarnConfiguration.NM_ADMIN_USER_ENV, defEnvStr, conf));
 
     // TODO: Remove Windows check and use this approach on all platforms after
     // additional testing.  See YARN-358.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunch.java

@@ -491,7 +491,6 @@ protected String getNMEnvVar(String varname) {
     }
   }
 
-
   @Test (timeout = 20000)
   public void testInvalidEnvSyntaxDiagnostics() throws IOException  {
 
@@ -688,9 +687,10 @@ public void handle(Event event) {
     resources.put(userjar, lpaths);
 
     Path nmp = new Path(testDir);
+    Set<String> nmEnvTrack = new LinkedHashSet<>();
 
     launch.sanitizeEnv(userSetEnv, pwd, appDirs, userLocalDirs, containerLogs,
-        resources, nmp, Collections.emptySet());
+        resources, nmp, nmEnvTrack);
 
     List<String> result =
       getJarManifestClasspath(userSetEnv.get(Environment.CLASSPATH.name()));
@@ -709,7 +709,7 @@ public void handle(Event event) {
         dispatcher, exec, null, container, dirsHandler, containerManager);
 
     launch.sanitizeEnv(userSetEnv, pwd, appDirs, userLocalDirs, containerLogs,
-        resources, nmp, Collections.emptySet());
+        resources, nmp, nmEnvTrack);
 
     result =
       getJarManifestClasspath(userSetEnv.get(Environment.CLASSPATH.name()));
@@ -720,6 +720,94 @@ public void handle(Event event) {
 
   }
 
+  @Test
+  public void testSanitizeNMEnvVars() throws Exception {
+    // Valid only for unix
+    assumeNotWindows();
+    ContainerLaunchContext containerLaunchContext =
+        recordFactory.newRecordInstance(ContainerLaunchContext.class);
+    ApplicationId appId = ApplicationId.newInstance(0, 0);
+    ApplicationAttemptId appAttemptId =
+        ApplicationAttemptId.newInstance(appId, 1);
+    ContainerId cId = ContainerId.newContainerId(appAttemptId, 0);
+    Map<String, String> userSetEnv = new HashMap<String, String>();
+    Set<String> nmEnvTrack = new LinkedHashSet<>();
+    userSetEnv.put(Environment.CONTAINER_ID.name(), "user_set_container_id");
+    userSetEnv.put(Environment.NM_HOST.name(), "user_set_NM_HOST");
+    userSetEnv.put(Environment.NM_PORT.name(), "user_set_NM_PORT");
+    userSetEnv.put(Environment.NM_HTTP_PORT.name(), "user_set_NM_HTTP_PORT");
+    userSetEnv.put(Environment.LOCAL_DIRS.name(), "user_set_LOCAL_DIR");
+    userSetEnv.put(Environment.USER.key(), "user_set_" +
+        Environment.USER.key());
+    userSetEnv.put(Environment.LOGNAME.name(), "user_set_LOGNAME");
+    userSetEnv.put(Environment.PWD.name(), "user_set_PWD");
+    userSetEnv.put(Environment.HOME.name(), "user_set_HOME");
+    userSetEnv.put(Environment.CLASSPATH.name(), "APATH");
+    // This one should be appended to.
+    String userMallocArenaMaxVal = "test_user_max_val";
+    userSetEnv.put("MALLOC_ARENA_MAX", userMallocArenaMaxVal);
+    containerLaunchContext.setEnvironment(userSetEnv);
+    Container container = mock(Container.class);
+    when(container.getContainerId()).thenReturn(cId);
+    when(container.getLaunchContext()).thenReturn(containerLaunchContext);
+    when(container.getLocalizedResources()).thenReturn(null);
+    Dispatcher dispatcher = mock(Dispatcher.class);
+    EventHandler<Event> eventHandler = new EventHandler<Event>() {
+      public void handle(Event event) {
+        Assert.assertTrue(event instanceof ContainerExitEvent);
+        ContainerExitEvent exitEvent = (ContainerExitEvent) event;
+        Assert.assertEquals(ContainerEventType.CONTAINER_EXITED_WITH_FAILURE,
+            exitEvent.getType());
+      }
+    };
+    when(dispatcher.getEventHandler()).thenReturn(eventHandler);
+
+    // these should eclipse anything in the user environment
+    YarnConfiguration conf = new YarnConfiguration();
+    String mallocArenaMaxVal = "test_nm_max_val";
+    conf.set("yarn.nodemanager.admin-env",
+        "MALLOC_ARENA_MAX=" + mallocArenaMaxVal);
+    String testKey1 = "TEST_KEY1";
+    String testVal1 = "testVal1";
+    conf.set("yarn.nodemanager.admin-env." + testKey1, testVal1);
+    String testKey2 = "TEST_KEY2";
+    String testVal2 = "testVal2";
+    conf.set("yarn.nodemanager.admin-env." + testKey2, testVal2);
+    String testKey3 = "MOUNT_LIST";
+    String testVal3 = "/home/a/b/c,/home/d/e/f,/home/g/e/h";
+    conf.set("yarn.nodemanager.admin-env." + testKey3, testVal3);
+    Map<String, String> environment = new HashMap<>();
+    LinkedHashSet<String> nmVars = new LinkedHashSet<>();
+    ContainerLaunch launch = new ContainerLaunch(distContext, conf,
+        dispatcher, exec, null, container, dirsHandler, containerManager);
+    String testDir = System.getProperty("test.build.data",
+        "target/test-dir");
+    Path pwd = new Path(testDir);
+    List<Path> appDirs = new ArrayList<Path>();
+    List<String> userLocalDirs = new ArrayList<>();
+    List<String> containerLogs = new ArrayList<String>();
+    Map<Path, List<String>> resources = new HashMap<Path, List<String>>();
+    Path userjar = new Path("user.jar");
+    List<String> lpaths = new ArrayList<String>();
+    lpaths.add("userjarlink.jar");
+    resources.put(userjar, lpaths);
+    Path nmp = new Path(testDir);
+
+    launch.sanitizeEnv(userSetEnv, pwd, appDirs, userLocalDirs, containerLogs,
+        resources, nmp, nmEnvTrack);
+    Assert.assertTrue(userSetEnv.containsKey("MALLOC_ARENA_MAX"));
+    Assert.assertTrue(userSetEnv.containsKey(testKey1));
+    Assert.assertTrue(userSetEnv.containsKey(testKey2));
+    Assert.assertTrue(userSetEnv.containsKey(testKey3));
+    Assert.assertTrue(nmEnvTrack.contains("MALLOC_ARENA_MAX"));
+    Assert.assertTrue(nmEnvTrack.contains("MOUNT_LIST"));
+    Assert.assertEquals(userMallocArenaMaxVal + File.pathSeparator
+        + mallocArenaMaxVal, userSetEnv.get("MALLOC_ARENA_MAX"));
+    Assert.assertEquals(testVal1, userSetEnv.get(testKey1));
+    Assert.assertEquals(testVal2, userSetEnv.get(testKey2));
+    Assert.assertEquals(testVal3, userSetEnv.get(testKey3));
+  }
+
   @Test
   public void testErrorLogOnContainerExit() throws Exception {
     verifyTailErrorLogOnContainerExit(new Configuration(), "/stderr", false);