From 0ba6f35dc24ace78ae666cf0a018b174d65bff17 Mon Sep 17 00:00:00 2001
From: Tamas Domok <domoktamas@gmail.com>
Date: Wed, 4 Aug 2021 15:27:46 +0200
Subject: [PATCH] YARN-10355. Refactor NM
 ContainerLaunch.java#orderEnvByDependencies

Co-authored-by: Tamas Domok <tdomok@cloudera.com>
---
 .../launcher/ContainerLaunch.java             | 48 +++++++++----------
 .../TestContainerLaunchParameterized.java     |  2 +
 2 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/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 b/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
index e864c14ad70..3673cee9e83 100644
--- a/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
+++ b/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
@@ -46,6 +46,8 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataInputStream;
@@ -1449,6 +1451,9 @@ public Set<String> getEnvDependencies(final String envVal) {
   private static final class WindowsShellScriptBuilder
       extends ShellScriptBuilder {
 
+    private static final Pattern VARIABLE_PATTERN = Pattern.compile("%(.*?)%");
+    private static final Pattern SPLIT_PATTERN = Pattern.compile(":");
+
     private void errorCheck() {
       line("@if %errorlevel% neq 0 exit /b %errorlevel%");
     }
@@ -1539,34 +1544,25 @@ public Set<String> getEnvDependencies(final String envVal) {
       if (envVal == null || envVal.isEmpty()) {
         return Collections.emptySet();
       }
+
+      // Example inputs: %var%, %%, %a:b%
+      Matcher matcher = VARIABLE_PATTERN.matcher(envVal);
       final Set<String> deps = new HashSet<>();
-      final int len = envVal.length();
-      int i = 0;
-      while (i < len) {
-        i = envVal.indexOf('%', i); // find beginning of variable
-        if (i < 0 || i == (len - 1)) {
-          break;
+      while (matcher.find()) {
+        String match = matcher.group(1);
+        if (!match.isEmpty()) {
+          if (match.equals(":")) {
+            // Special case, variable name can be a single : character
+            deps.add(match);
+          } else {
+            // Either store the variable name before the : string manipulation
+            // character or the whole match. (%var% -> var, %a:b% -> a)
+            String[] split = SPLIT_PATTERN.split(match, 2);
+            if (!split[0].isEmpty()) {
+              deps.add(split[0]);
+            }
+          }
         }
-        i++;
-        // 3 cases: %var%, %var:...% or %%
-        final int j = envVal.indexOf('%', i); // find end of variable
-        if (j == i) {
-          // %% case, just skip it
-          i++;
-          continue;
-        }
-        if (j < 0) {
-          break; // even %var:...% syntax ends with a %, so j cannot be negative
-        }
-        final int k = envVal.indexOf(':', i);
-        if (k >= 0 && k < j) {
-          // %var:...% syntax
-          deps.add(envVal.substring(i, k));
-        } else {
-          // %var% syntax
-          deps.add(envVal.substring(i, j));
-        }
-        i = j + 1;
       }
       return deps;
     }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunchParameterized.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunchParameterized.java
index 17a5825289a..a36d5efc2c3 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunchParameterized.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunchParameterized.java
@@ -118,6 +118,8 @@ private static Stream<Arguments> inputForGetEnvDependenciesWin() {
         Arguments.of("%A", asSet()),
         Arguments.of("%A:", asSet()),
         Arguments.of("%A%", asSet("A")),
+        Arguments.of("%:%", asSet(":")),
+        Arguments.of("%:A%", asSet()),
         Arguments.of("%%%A%", asSet("A")),
         Arguments.of("%%C%A%", asSet("A")),
         Arguments.of("%A:~-1%", asSet("A")),