YARN-10664. Allow parameter expansion in NM_ADMIN_USER_ENV. Contributed by Jim

Brennan.

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

@@ -175,15 +175,13 @@ public class ContainerLaunch implements Callable<Integer> {
     return var;
   }
 
-  private Map<String, String> expandAllEnvironmentVars(
-      ContainerLaunchContext launchContext, Path containerLogDir) {
-    Map<String, String> environment = launchContext.getEnvironment();
+  private void expandAllEnvironmentVars(
+      Map<String, String> environment, Path containerLogDir) {
     for (Entry<String, String> entry : environment.entrySet()) {
       String value = entry.getValue();
       value = expandEnvironment(value, containerLogDir);
       entry.setValue(value);
     }
-    return environment;
   }
 
   @Override
@@ -218,8 +216,10 @@ public class ContainerLaunch implements Callable<Integer> {
       }
       launchContext.setCommands(newCmds);
 
-      Map<String, String> environment = expandAllEnvironmentVars(
-          launchContext, containerLogDir);
+      // The actual expansion of environment variables happens after calling
+      // sanitizeEnv.  This allows variables specified in NM_ADMIN_USER_ENV
+      // to reference user or container-defined variables.
+      Map<String, String> environment = launchContext.getEnvironment();
       // /////////////////////////// End of variable expansion
 
       // Use this to track variables that are added to the environment by nm.
@@ -283,6 +283,8 @@ public class ContainerLaunch implements Callable<Integer> {
             containerLogDirs, localResources, nmPrivateClasspathJarDir,
             nmEnvVars);
 
+        expandAllEnvironmentVars(environment, containerLogDir);
+
         prepareContainer(localResources, containerLocalDirs);
 
         // Write out the environment
@@ -1564,13 +1566,13 @@ public class ContainerLaunch implements Callable<Integer> {
     }
 
     // variables here will be forced in, even if the container has
-    // specified them.
+    // specified them.  Note: we do not track these in nmVars, to
+    // allow them to be ordered properly if they reference variables
+    // defined by the user.
     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));
 
     if (!Shell.WINDOWS) {
       // maybe force path components

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

@@ -25,6 +25,7 @@ import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.*;
 
 import java.io.BufferedReader;
+import java.io.DataOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.FileReader;
@@ -55,9 +56,8 @@ import java.util.function.Supplier;
 import com.google.common.collect.Lists;
 import org.apache.commons.codec.binary.Base64;
 import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.fs.FileUtil;
-import org.apache.hadoop.fs.Path;
-import org.apache.hadoop.fs.UnsupportedFileSystemException;
+import org.apache.hadoop.fs.*;
+import org.apache.hadoop.io.IOUtils;
 import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.security.token.SecretManager.InvalidToken;
 import org.apache.hadoop.test.GenericTestUtils;
@@ -124,6 +124,8 @@ import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 
 public class TestContainerLaunch extends BaseContainerManagerTest {
 
@@ -664,7 +666,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     Container container = mock(Container.class);
     when(container.getContainerId()).thenReturn(cId);
     when(container.getLaunchContext()).thenReturn(containerLaunchContext);
-    when(container.getLocalizedResources()).thenReturn(null);
+    when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5");
     Dispatcher dispatcher = mock(Dispatcher.class);
     EventHandler<Event> eventHandler = new EventHandler<Event>() {
       public void handle(Event event) {
@@ -805,8 +807,6 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     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));
@@ -1848,6 +1848,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     when(id.toString()).thenReturn("1");
     when(container.getContainerId()).thenReturn(id);
     when(container.getUser()).thenReturn("user");
+    when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5");
     ContainerLaunchContext clc = mock(ContainerLaunchContext.class);
     when(clc.getCommands()).thenReturn(Lists.newArrayList());
     when(container.getLaunchContext()).thenReturn(clc);
@@ -2444,6 +2445,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
         .newContainerId(ApplicationAttemptId.newInstance(appId, 1), 1);
     when(container.getContainerId()).thenReturn(containerId);
     when(container.getUser()).thenReturn("test");
+    when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5");
 
     when(container.getLocalizedResources())
         .thenReturn(Collections.<Path, List<String>> emptyMap());
@@ -2512,4 +2514,102 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
         launch.getUserFilecacheDirs(localDirsForRead)),
         StringUtils.join(",", ctx.getUserFilecacheDirs()));
   }
+
+  private String readStringFromPath(Path p) throws IOException {
+    FileSystem fs = FileSystem.get(conf);
+    try (FSDataInputStream is = fs.open(p)) {
+      byte[] bytes = IOUtils.readFullyToByteArray(is);
+      return new String(bytes);
+    }
+  }
+
+  @Test(timeout = 20000)
+  public void testExpandNmAdmEnv() throws Exception {
+    // setup mocks
+    Dispatcher dispatcher = mock(Dispatcher.class);
+    EventHandler handler = mock(EventHandler.class);
+    when(dispatcher.getEventHandler()).thenReturn(handler);
+    ContainerExecutor containerExecutor = mock(ContainerExecutor.class);
+    doAnswer(new Answer<Void>() {
+      @Override
+      public Void answer(InvocationOnMock invocation) throws Throwable {
+        Object[] args = invocation.getArguments();
+        DataOutputStream dos = (DataOutputStream) args[0];
+        dos.writeBytes("script");
+        return null;
+      }
+    }).when(containerExecutor).writeLaunchEnv(
+        any(), any(), any(), any(), any(), any(), any());
+    Application app = mock(Application.class);
+    ApplicationId appId = mock(ApplicationId.class);
+    when(appId.toString()).thenReturn("1");
+    when(app.getAppId()).thenReturn(appId);
+    Container container = mock(Container.class);
+    ContainerId id = mock(ContainerId.class);
+    when(id.toString()).thenReturn("1");
+    when(container.getContainerId()).thenReturn(id);
+    when(container.getUser()).thenReturn("user");
+    ContainerLaunchContext clc = mock(ContainerLaunchContext.class);
+    when(clc.getCommands()).thenReturn(Lists.newArrayList());
+    when(container.getLaunchContext()).thenReturn(clc);
+    Credentials credentials = mock(Credentials.class);
+    when(container.getCredentials()).thenReturn(credentials);
+    when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5");
+
+    // Define user environment variables.
+    Map<String, String> userSetEnv = new HashMap<String, String>();
+    String userVar = "USER_VAR";
+    String userVarVal = "user-var-value";
+    userSetEnv.put(userVar, userVarVal);
+    when(clc.getEnvironment()).thenReturn(userSetEnv);
+
+    YarnConfiguration localConf = new YarnConfiguration(conf);
+
+    // Admin Env var that depends on USER_VAR1
+    String testKey1 = "TEST_KEY1";
+    String testVal1 = "relies on {{USER_VAR}}";
+    localConf.set(
+        YarnConfiguration.NM_ADMIN_USER_ENV + "." + testKey1, testVal1);
+    String testVal1Expanded; // this is what we expect after {{}} expansion
+    if (Shell.WINDOWS) {
+      testVal1Expanded = "relies on %USER_VAR%";
+    } else {
+      testVal1Expanded = "relies on $USER_VAR";
+    }
+    // Another Admin Env var that depends on the first one
+    String testKey2 = "TEST_KEY2";
+    String testVal2 = "relies on {{TEST_KEY1}}";
+    localConf.set(
+        YarnConfiguration.NM_ADMIN_USER_ENV + "." + testKey2, testVal2);
+    String testVal2Expanded; // this is what we expect after {{}} expansion
+    if (Shell.WINDOWS) {
+      testVal2Expanded = "relies on %TEST_KEY1%";
+    } else {
+      testVal2Expanded = "relies on $TEST_KEY1";
+    }
+
+    // call containerLaunch
+    ContainerLaunch containerLaunch = new ContainerLaunch(
+        distContext, localConf, dispatcher,
+        containerExecutor, app, container, dirsHandler, containerManager);
+    containerLaunch.call();
+
+    // verify the nmPrivate paths and files
+    ArgumentCaptor<ContainerStartContext> cscArgument =
+        ArgumentCaptor.forClass(ContainerStartContext.class);
+    verify(containerExecutor, times(1)).launchContainer(cscArgument.capture());
+    ContainerStartContext csc = cscArgument.getValue();
+    Assert.assertEquals("script",
+        readStringFromPath(csc.getNmPrivateContainerScriptPath()));
+
+    // verify env
+    ArgumentCaptor<Map> envArgument = ArgumentCaptor.forClass(Map.class);
+    verify(containerExecutor, times(1)).writeLaunchEnv(any(),
+        envArgument.capture(), any(), any(), any(), any(), any());
+    Map env = envArgument.getValue();
+    Assert.assertEquals(userVarVal, env.get(userVar));
+    Assert.assertEquals(testVal1Expanded, env.get(testKey1));
+    Assert.assertEquals(testVal2Expanded, env.get(testKey2));
+  }
+
 }