Added run_under_env argument to RunEnvironmentInfo

src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java

@@ -474,6 +474,7 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
     if (environmentProvider != null) {
       testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
       testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());
+      testActionBuilder.setRunUnderEnv(environmentProvider.getRunUnderEnv());
     }
 
     TestParams testParams =

src/main/java/com/google/devtools/build/lib/analysis/RunEnvironmentInfo.java

@@ -23,9 +23,11 @@
 import com.google.devtools.build.lib.starlarkbuildapi.RunEnvironmentInfoApi;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Sequence;
+import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkList;
 
 /**
@@ -40,16 +42,19 @@ public final class RunEnvironmentInfo extends NativeInfo implements RunEnvironme
 
   private final ImmutableMap<String, String> environment;
   private final ImmutableList<String> inheritedEnvironment;
+  private final @Nullable String runUnderEnv;
   private final boolean shouldErrorOnNonExecutableRule;
 
   /** Constructs a new provider with the given fixed and inherited environment variables. */
   public RunEnvironmentInfo(
       Map<String, String> environment,
       List<String> inheritedEnvironment,
+      @Nullable String runUnderEnv,
       boolean shouldErrorOnNonExecutableRule) {
     this.environment = ImmutableMap.copyOf(Preconditions.checkNotNull(environment));
     this.inheritedEnvironment =
         ImmutableList.copyOf(Preconditions.checkNotNull(inheritedEnvironment));
+    this.runUnderEnv = runUnderEnv;
     this.shouldErrorOnNonExecutableRule = shouldErrorOnNonExecutableRule;
   }
 
@@ -76,6 +81,12 @@ public ImmutableList<String> getInheritedEnvironment() {
     return inheritedEnvironment;
   }
 
+  /** Returns the name of the environment variable to pass the "run_under" value to. */
+  @Nullable
+  public String getRunUnderEnv() {
+    return runUnderEnv;
+  }
+
   /**
    * Returns whether advertising this provider on a non-executable (and thus non-test) rule should
    * result in an error or a warning. The latter is required to not break testing.TestEnvironment,
@@ -95,11 +106,13 @@ private RunEnvironmentInfoProvider() {
 
     @Override
     public RunEnvironmentInfoApi constructor(
-        Dict<?, ?> environment, Sequence<?> inheritedEnvironment) throws EvalException {
+        Dict<?, ?> environment, Sequence<?> inheritedEnvironment, Object runUnderEnv)
+        throws EvalException {
       return new RunEnvironmentInfo(
           Dict.cast(environment, String.class, String.class, "environment"),
           StarlarkList.immutableCopyOf(
               Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")),
+          runUnderEnv == Starlark.NONE ? null : (String) runUnderEnv,
           /* shouldErrorOnNonExecutableRule= */ true);
     }
   }

src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java

@@ -98,12 +98,14 @@ public NestedSet<PackageGroupContents> getPackageSpecifications() {
   private int explicitShardCount;
   private final Map<String, String> extraEnv;
   private final Set<String> extraInheritedEnv;
+  private @Nullable String runUnderEnv;
 
   public TestActionBuilder(RuleContext ruleContext) {
     this.ruleContext = ruleContext;
     this.extraEnv = new TreeMap<>();
     this.extraInheritedEnv = new TreeSet<>();
     this.additionalTools = new ImmutableList.Builder<>();
+    this.runUnderEnv = null;
   }
 
   /**
@@ -180,6 +182,12 @@ public TestActionBuilder addExtraInheritedEnv(List<String> extraInheritedEnv) {
     return this;
   }
 
+  @CanIgnoreReturnValue
+  public TestActionBuilder setRunUnderEnv(@Nullable String runUnderEnv) {
+    this.runUnderEnv = runUnderEnv;
+    return this;
+  }
+
   /** Set the explicit shard count. Note that this may be overridden by the sharding strategy. */
   @CanIgnoreReturnValue
   public TestActionBuilder setShardCount(int explicitShardCount) {
@@ -340,6 +348,7 @@ private TestParams createTestAction(int shards)
               ruleContext,
               runfilesSupport,
               executable,
+              runUnderEnv,
               instrumentedFileManifest,
               shards,
               runsPerTest);
@@ -352,7 +361,7 @@ private TestParams createTestAction(int shards)
     } else {
       executionSettings =
           new TestTargetExecutionSettings(
-              ruleContext, runfilesSupport, executable, null, shards, runsPerTest);
+              ruleContext, runfilesSupport, executable, runUnderEnv, null, shards, runsPerTest);
     }
 
     extraTestEnv.putAll(extraEnv);

src/main/java/com/google/devtools/build/lib/analysis/test/TestStrategy.java

@@ -179,9 +179,10 @@ public static ImmutableList<String> expandedArgsFromAction(TestRunnerAction test
     }
 
     TestTargetExecutionSettings execSettings = testAction.getExecutionSettings();
+    boolean includeRunUnderArgs = execSettings.getRunUnderEnv() == null;
 
     // Insert the command prefix specified by the "--run_under=<command-prefix>" option, if any.
-    if (execSettings.getRunUnder() != null) {
+    if (includeRunUnderArgs && execSettings.getRunUnder() != null) {
       addRunUnderArgs(testAction, args, executedOnWindows);
     }
 
@@ -191,25 +192,42 @@ public static ImmutableList<String> expandedArgsFromAction(TestRunnerAction test
     return ImmutableList.copyOf(args);
   }
 
-  private static void addRunUnderArgs(
-      TestRunnerAction testAction, List<String> args, boolean executedOnWindows) {
+  /**
+   * Get {@code --run_under} command line values.
+   *
+   * @param testAction The test action.
+   * @return the {@code --run_under} arguments as string list.
+   * @throws CommandLineExpansionException
+   */
+  public static List<String> computeRunUnderArgs(TestRunnerAction testAction) {
+    List<String> args = Lists.newArrayList();
     TestTargetExecutionSettings execSettings = testAction.getExecutionSettings();
     if (execSettings.getRunUnderExecutable() != null) {
       args.add(execSettings.getRunUnderExecutable().getRunfilesPath().getCallablePathString());
     } else {
-      if (execSettings.needsShell(executedOnWindows)) {
-        // TestActionBuilder constructs TestRunnerAction with a 'null' shell only when none is
-        // required. Something clearly went wrong.
-        Preconditions.checkNotNull(testAction.getShExecutableMaybe(), "%s", testAction);
-        String shellExecutable = testAction.getShExecutableMaybe().getPathString();
-        args.add(shellExecutable);
-        args.add("-c");
-        args.add("\"$@\"");
-        args.add(shellExecutable); // Sets $0.
-      }
       args.add(execSettings.getRunUnder().getCommand());
     }
     args.addAll(testAction.getExecutionSettings().getRunUnder().getOptions());
+
+    return args;
+  }
+
+  private static void addRunUnderArgs(
+      TestRunnerAction testAction, List<String> args, boolean executedOnWindows) {
+    TestTargetExecutionSettings execSettings = testAction.getExecutionSettings();
+    if (execSettings.getRunUnderExecutable() == null
+        && execSettings.needsShell(executedOnWindows)) {
+      // TestActionBuilder constructs TestRunnerAction with a 'null' shell only when none is
+      // required. Something clearly went wrong.
+      Preconditions.checkNotNull(testAction.getShExecutableMaybe(), "%s", testAction);
+      String shellExecutable = testAction.getShExecutableMaybe().getPathString();
+      args.add(shellExecutable);
+      args.add("-c");
+      args.add("\"$@\"");
+      args.add(shellExecutable); // Sets $0.
+    }
+
+    args.addAll(computeRunUnderArgs(testAction));
   }
 
   /**

src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetExecutionSettings.java

@@ -41,6 +41,7 @@ public final class TestTargetExecutionSettings {
   private final int totalRuns;
   private final RunUnder runUnder;
   private final Artifact runUnderExecutable;
+  private final @Nullable String runUnderEnv;
   private final Artifact executable;
   private final boolean runfilesSymlinksCreated;
   @Nullable private final Path runfilesDir;
@@ -53,6 +54,7 @@ public final class TestTargetExecutionSettings {
       RuleContext ruleContext,
       RunfilesSupport runfilesSupport,
       Artifact executable,
+      @Nullable String runUnderEnv,
       Artifact instrumentedFileManifest,
       int shards,
       int runs)
@@ -74,6 +76,7 @@ public final class TestTargetExecutionSettings {
     this.testFilter = testConfig.getTestFilter();
     this.testRunnerFailFast = testConfig.getTestRunnerFailFast();
     this.executable = executable;
+    this.runUnderEnv = runUnderEnv;
     this.runfilesSymlinksCreated = runfilesSupport.isBuildRunfileLinks();
     this.runfilesDir = runfilesSupport.getRunfilesDirectory();
     this.runfiles = runfilesSupport.getRunfiles();
@@ -117,6 +120,10 @@ public RunUnder getRunUnder() {
     return runUnder;
   }
 
+  public @Nullable String getRunUnderEnv() {
+    return runUnderEnv;
+  }
+
   public Artifact getExecutable() {
     return executable;
   }

src/main/java/com/google/devtools/build/lib/exec/BUILD

@@ -438,8 +438,10 @@ java_library(
     name = "test_policy",
     srcs = ["TestPolicy.java"],
     deps = [
+        "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/util:shell_escaper",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:guava",
     ],

src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java

@@ -15,6 +15,8 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.test.TestRunnerAction;
+import com.google.devtools.build.lib.analysis.test.TestStrategy;
+import com.google.devtools.build.lib.util.ShellEscaper;
 import com.google.devtools.build.lib.util.UserUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.time.Duration;
@@ -90,6 +92,13 @@ public Map<String, String> computeTestEnvironment(
     // Rule-specified test env.
     testAction.getExtraTestEnv().resolve(env, clientEnv);
 
+    // Add the --run_under environment variable if set.
+    String runUnderEnv = testAction.getExecutionSettings().getRunUnderEnv();
+    if (runUnderEnv != null) {
+      env.put(
+          runUnderEnv, ShellEscaper.escapeJoinAll(TestStrategy.computeRunUnderArgs(testAction)));
+    }
+
     // Overwrite with the environment common to all actions, see --action_env.
     testAction.getConfiguration().getActionEnvironment().resolve(env, clientEnv);
 

src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java

@@ -49,6 +49,7 @@ public RunEnvironmentInfo testEnvironment(
         Dict.cast(environment, String.class, String.class, "environment"),
         StarlarkList.immutableCopyOf(
             Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")),
+        /* runUnderEnv= */ null,
         /* shouldErrorOnNonExecutableRule= */ false);
   }
 

src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java

@@ -194,6 +194,30 @@ private List<String> computeArgs(ConfiguredTarget targetToRun, List<String> comm
     return args;
   }
 
+  private String computeRunUnderArgs(
+      RunUnder runUnder, ConfiguredTarget targetToRun, ConfiguredTarget runUnderTarget) {
+    String runUnderValue = runUnder.getValue();
+    if (runUnderTarget != null) {
+      // --run_under specifies a target. Get the corresponding executable.
+      // This must be an absolute path, because the run_under target is only
+      // in the runfiles of test targets.
+      runUnderValue =
+          runUnderTarget
+              .getProvider(FilesToRunProvider.class)
+              .getExecutable()
+              .getPath()
+              .getPathString();
+      // If the run_under command contains any options, make sure to add them
+      // to the command line as well.
+      List<String> opts = runUnder.getOptions();
+      if (!opts.isEmpty()) {
+        runUnderValue += " " + ShellEscaper.escapeJoinAll(opts);
+      }
+    }
+
+    return runUnderValue;
+  }
+
   private void constructCommandLine(
       List<String> cmdLine,
       List<String> prettyCmdLine,
@@ -222,24 +246,16 @@ private void constructCommandLine(
     PathFragment prettyExecutablePath =
         prettyPrinter.getPrettyPath(executable.getPath().asFragment());
 
+    RunEnvironmentInfo runEnvironmentInfo =
+        (RunEnvironmentInfo) targetToRun.get(RunEnvironmentInfo.PROVIDER.getKey());
+    boolean runUnderEnvDefined =
+        (runEnvironmentInfo != null && runEnvironmentInfo.getRunUnderEnv() != null);
+
     RunUnder runUnder = env.getOptions().getOptions(CoreOptions.class).runUnder;
-    // Insert the command prefix specified by the "--run_under=<command-prefix>" option
-    // at the start of the command line.
-    if (runUnder != null) {
-      String runUnderValue = runUnder.getValue();
-      if (runUnderTarget != null) {
-        // --run_under specifies a target. Get the corresponding executable.
-        // This must be an absolute path, because the run_under target is only
-        // in the runfiles of test targets.
-        runUnderValue = runUnderTarget
-            .getProvider(FilesToRunProvider.class).getExecutable().getPath().getPathString();
-        // If the run_under command contains any options, make sure to add them
-        // to the command line as well.
-        List<String> opts = runUnder.getOptions();
-        if (!opts.isEmpty()) {
-          runUnderValue += " " + ShellEscaper.escapeJoinAll(opts);
-        }
-      }
+    // Insert the command prefix specified by the "--run_under=<command-prefix>" option at
+    // the start of the command line whenever RunEnvironmentInfo.run_under_env is not used.
+    if (runUnder != null && !runUnderEnvDefined) {
+      String runUnderValue = computeRunUnderArgs(runUnder, targetToRun, runUnderTarget);
 
       PathFragment shellExecutable = ShToolchain.getPathForHost(configuration);
       if (shellExecutable.isEmpty()) {
@@ -500,6 +516,14 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
             actionEnvironment.withAdditionalVariables(
                 environmentProvider.getEnvironment(),
                 ImmutableSet.copyOf(environmentProvider.getInheritedEnvironment()));
+        if (environmentProvider.getRunUnderEnv() != null) {
+          actionEnvironment =
+              actionEnvironment.withAdditionalVariables(
+                  ImmutableMap.of(
+                      environmentProvider.getRunUnderEnv(),
+                      computeRunUnderArgs(runUnder, targetToRun, runUnderTarget)),
+                  /* inheritedVars= */ ImmutableSet.of());
+        }
       }
       actionEnvironment.resolve(runEnvironment, env.getClientEnv());
       try {

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/RunEnvironmentInfoApi.java

@@ -20,12 +20,14 @@
 import com.google.devtools.build.lib.starlarkbuildapi.core.StructApi;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 import net.starlark.java.annot.Param;
 import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.NoneType;
 import net.starlark.java.eval.Sequence;
 
 /**
@@ -91,13 +93,25 @@ interface RunEnvironmentInfoApiProvider extends ProviderApi {
                       + " when the target that returns this provider is executed, either as a"
                       + " test or via the run command. If a variable is contained in both <code>"
                       + "environment</code> and <code>inherited_environment</code>, the value"
-                      + " inherited from the shell environment will take precedence if set.")
+                      + " inherited from the shell environment will take precedence if set."),
+          @Param(
+              name = "run_under_env",
+              allowedTypes = {@ParamType(type = String.class), @ParamType(type = NoneType.class)},
+              defaultValue = "None",
+              named = true,
+              positional = false,
+              doc =
+                  "The name of an environment variable to represent the value of the"
+                      + " <code>--run_under<code> command line argument. By setting this, the \"run"
+                      + " under\" functionality will not be executed. Instead, it will be up to the"
+                      + " built executable to handle this functionality.")
         },
         selfCall = true)
     @StarlarkConstructor
     RunEnvironmentInfoApi constructor(
         Dict<?, ?> environment, // <String, String> expected
-        Sequence<?> inheritedEnvironment /* <String> expected */)
+        Sequence<?> inheritedEnvironment, /* <String> expected */
+        Object run_under_env /* String or NoneType expected */)
         throws EvalException;
   }
 }