Let Starlark tests inherit env variables

Adds an inherited_environment field to testing.TestEnvironment to allow
Starlark rules to implement the equivalent of the native rules'
`env_inherit` attribute.

Work towards bazelbuild#7364. To fully resolve that issue, it remains to handle
executable non-test rules.

RELNOTES: Starlark test rules can use the new inherited_environment
parameter of testing.TestEnvironment to specify environment variables
whose values should be inherited from the shell environment.

Closes bazelbuild#14849.

PiperOrigin-RevId: 439277689

Author: fmeum

src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java

@@ -701,7 +701,7 @@ public Dict<String, String> getExecutionInfoDict() {
 
   @Override
   public Dict<String, String> getEnv() {
-    return Dict.immutableCopyOf(env.getFixedEnv().toMap());
+    return Dict.immutableCopyOf(env.getFixedEnv());
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java

@@ -18,7 +18,6 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.util.Fingerprint;
-import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -45,23 +44,34 @@
  */
 public final class ActionEnvironment {
 
-  /** A map of environment variables. */
+  /**
+   * A map of environment variables together with a list of variables to inherit from the shell
+   * environment.
+   */
   public interface EnvironmentVariables {
 
     /**
-     * Returns the environment variables as a map.
+     * Returns the fixed environment variables as a map.
      *
-     * <p>WARNING: this allocations additional objects if the underlying implementation is a {@link
+     * <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
      * CompoundEnvironmentVariables}; use sparingly.
      */
-    ImmutableMap<String, String> toMap();
+    ImmutableMap<String, String> getFixedEnvironment();
+
+    /**
+     * Returns the inherited environment variables as a set.
+     *
+     * <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
+     * CompoundEnvironmentVariables}; use sparingly.
+     */
+    ImmutableSet<String> getInheritedEnvironment();
 
     default boolean isEmpty() {
-      return toMap().isEmpty();
+      return getFixedEnvironment().isEmpty() && getInheritedEnvironment().isEmpty();
     }
 
     default int size() {
-      return toMap().size();
+      return getFixedEnvironment().size() + getInheritedEnvironment().size();
     }
   }
 
@@ -70,11 +80,21 @@ default int size() {
    * allocation a new map.
    */
   static class CompoundEnvironmentVariables implements EnvironmentVariables {
+
+    static EnvironmentVariables create(
+        Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
+      if (fixedVars.isEmpty() && inheritedVars.isEmpty() && base.isEmpty()) {
+        return EMPTY_ENVIRONMENT_VARIABLES;
+      }
+      return new CompoundEnvironmentVariables(fixedVars, inheritedVars, base);
+    }
+
     private final EnvironmentVariables current;
     private final EnvironmentVariables base;
 
-    CompoundEnvironmentVariables(Map<String, String> vars, EnvironmentVariables base) {
-      this.current = new SimpleEnvironmentVariables(vars);
+    private CompoundEnvironmentVariables(
+        Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
+      this.current = SimpleEnvironmentVariables.create(fixedVars, inheritedVars);
       this.base = base;
     }
 
@@ -84,48 +104,62 @@ public boolean isEmpty() {
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
-      Map<String, String> result = new LinkedHashMap<>();
-      result.putAll(base.toMap());
-      result.putAll(current.toMap());
-      return ImmutableMap.copyOf(result);
+    public ImmutableMap<String, String> getFixedEnvironment() {
+      ImmutableMap.Builder<String, String> result = new ImmutableMap.Builder<>();
+      result.putAll(base.getFixedEnvironment());
+      result.putAll(current.getFixedEnvironment());
+      return result.buildKeepingLast();
+    }
+
+    @Override
+    public ImmutableSet<String> getInheritedEnvironment() {
+      ImmutableSet.Builder<String> result = new ImmutableSet.Builder<>();
+      result.addAll(base.getInheritedEnvironment());
+      result.addAll(current.getInheritedEnvironment());
+      return result.build();
     }
   }
 
   /** A simple {@link EnvironmentVariables}. */
   static class SimpleEnvironmentVariables implements EnvironmentVariables {
 
-    static EnvironmentVariables create(Map<String, String> vars) {
-      if (vars.isEmpty()) {
+    static EnvironmentVariables create(Map<String, String> fixedVars, Set<String> inheritedVars) {
+      if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
         return EMPTY_ENVIRONMENT_VARIABLES;
       }
-      return new SimpleEnvironmentVariables(vars);
+      return new SimpleEnvironmentVariables(fixedVars, inheritedVars);
     }
 
-    private final ImmutableMap<String, String> vars;
+    private final ImmutableMap<String, String> fixedVars;
+    private final ImmutableSet<String> inheritedVars;
 
-    private SimpleEnvironmentVariables(Map<String, String> vars) {
-      this.vars = ImmutableMap.copyOf(vars);
+    private SimpleEnvironmentVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
+      this.fixedVars = ImmutableMap.copyOf(fixedVars);
+      this.inheritedVars = ImmutableSet.copyOf(inheritedVars);
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
-      return vars;
+    public ImmutableMap<String, String> getFixedEnvironment() {
+      return fixedVars;
+    }
+
+    @Override
+    public ImmutableSet<String> getInheritedEnvironment() {
+      return inheritedVars;
     }
   }
 
   /** An empty {@link EnvironmentVariables}. */
   public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
-      new SimpleEnvironmentVariables(ImmutableMap.of());
+      new SimpleEnvironmentVariables(ImmutableMap.of(), ImmutableSet.of());
 
   /**
    * An empty environment, mainly for testing. Production code should never use this, but instead
    * get the proper environment from the current configuration.
    */
   // TODO(ulfjack): Migrate all production code to use the proper action environment, and then make
   // this @VisibleForTesting or rename it to clarify.
-  public static final ActionEnvironment EMPTY =
-      new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES, ImmutableSet.of());
+  public static final ActionEnvironment EMPTY = new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES);
 
   /**
    * Splits the given map into a map of variables with a fixed value, and a set of variables that
@@ -145,59 +179,66 @@ public static ActionEnvironment split(Map<String, String> env) {
         inheritedEnv.add(key);
       }
     }
-    return create(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
+    return create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
   }
 
-  private final EnvironmentVariables fixedEnv;
-  private final ImmutableSet<String> inheritedEnv;
+  private final EnvironmentVariables vars;
 
-  private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
-    this.fixedEnv = fixedEnv;
-    this.inheritedEnv = inheritedEnv;
+  private ActionEnvironment(EnvironmentVariables vars) {
+    this.vars = vars;
   }
 
   /**
    * Creates a new action environment. The order in which the environments are combined is
    * undefined, so callers need to take care that the key set of the {@code fixedEnv} map and the
    * set of {@code inheritedEnv} elements are disjoint.
    */
-  private static ActionEnvironment create(
-      EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
-    if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
+  private static ActionEnvironment create(EnvironmentVariables vars) {
+    if (vars.isEmpty()) {
       return EMPTY;
     }
-    return new ActionEnvironment(fixedEnv, inheritedEnv);
+    return new ActionEnvironment(vars);
   }
 
   public static ActionEnvironment create(
       Map<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
-    return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv);
+    return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
   }
 
   public static ActionEnvironment create(Map<String, String> fixedEnv) {
-    return new ActionEnvironment(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.of());
+    return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, ImmutableSet.of()));
   }
 
   /**
    * Returns a copy of the environment with the given fixed variables added to it, <em>overwriting
    * any existing occurrences of those variables</em>.
    */
-  public ActionEnvironment addFixedVariables(Map<String, String> vars) {
-    return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv);
+  public ActionEnvironment addFixedVariables(Map<String, String> fixedVars) {
+    return ActionEnvironment.create(
+        CompoundEnvironmentVariables.create(fixedVars, ImmutableSet.of(), vars));
+  }
+
+  /**
+   * Returns a copy of the environment with the given fixed and inherited variables added to it,
+   * <em>overwriting any existing occurrences of those variables</em>.
+   */
+  public ActionEnvironment addVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
+    return ActionEnvironment.create(
+        CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars));
   }
 
   /** Returns the combined size of the fixed and inherited environments. */
   public int size() {
-    return fixedEnv.size() + inheritedEnv.size();
+    return vars.size();
   }
 
   /**
    * Returns the 'fixed' part of the environment, i.e., those environment variables that are set to
    * fixed values and their values. This should only be used for testing and to compute the cache
    * keys of actions. Use {@link #resolve} instead to get the complete environment.
    */
-  public EnvironmentVariables getFixedEnv() {
-    return fixedEnv;
+  public ImmutableMap<String, String> getFixedEnv() {
+    return vars.getFixedEnvironment();
   }
 
   /**
@@ -207,7 +248,7 @@ public EnvironmentVariables getFixedEnv() {
    * get the complete environment.
    */
   public ImmutableSet<String> getInheritedEnv() {
-    return inheritedEnv;
+    return vars.getInheritedEnvironment();
   }
 
   /**
@@ -218,8 +259,8 @@ public ImmutableSet<String> getInheritedEnv() {
    */
   public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
     checkNotNull(clientEnv);
-    result.putAll(fixedEnv.toMap());
-    for (String var : inheritedEnv) {
+    result.putAll(vars.getFixedEnvironment());
+    for (String var : vars.getInheritedEnvironment()) {
       String value = clientEnv.get(var);
       if (value != null) {
         result.put(var, value);
@@ -228,7 +269,7 @@ public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
   }
 
   public void addTo(Fingerprint f) {
-    f.addStringMap(fixedEnv.toMap());
-    f.addStrings(inheritedEnv);
+    f.addStringMap(vars.getFixedEnvironment());
+    f.addStrings(vars.getInheritedEnvironment());
   }
 }

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

@@ -470,6 +470,7 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
             providersBuilder.getProvider(TestEnvironmentInfo.PROVIDER.getKey());
     if (environmentProvider != null) {
       testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
+      testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());
     }
 
     TestParams testParams =

src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java

@@ -442,7 +442,7 @@ public String describeKey() {
     StringBuilder message = new StringBuilder();
     message.append(getProgressMessage());
     message.append('\n');
-    for (Map.Entry<String, String> entry : env.getFixedEnv().toMap().entrySet()) {
+    for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
       message.append("  Environment variable: ");
       message.append(ShellEscaper.escapeString(entry.getKey()));
       message.append('=');
@@ -551,7 +551,7 @@ public final ImmutableMap<String, String> getIncompleteEnvironmentForTesting() {
     // ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
     // requires first updating all subclasses and callers to actually handle environments correctly,
     // so it's not a small change.
-    return env.getFixedEnv().toMap();
+    return env.getFixedEnv();
   }
 
   /**

src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java

@@ -446,7 +446,7 @@ public boolean isSiblingRepositoryLayout() {
    */
   @Override
   public ImmutableMap<String, String> getLocalShellEnvironment() {
-    return actionEnv.getFixedEnv().toMap();
+    return actionEnv.getFixedEnv();
   }
 
   /**
@@ -583,7 +583,7 @@ public boolean legacyExternalRunfiles() {
    */
   @Override
   public ImmutableMap<String, String> getTestEnv() {
-    return testEnv.getFixedEnv().toMap();
+    return testEnv.getFixedEnv();
   }
 
   /**

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

@@ -55,7 +55,9 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.TreeMap;
+import java.util.TreeSet;
 import javax.annotation.Nullable;
 
 /**
@@ -96,10 +98,12 @@ public NestedSet<PackageGroupContents> getPackageSpecifications() {
   private InstrumentedFilesInfo instrumentedFiles;
   private int explicitShardCount;
   private final Map<String, String> extraEnv;
+  private final Set<String> extraInheritedEnv;
 
   public TestActionBuilder(RuleContext ruleContext) {
     this.ruleContext = ruleContext;
     this.extraEnv = new TreeMap<>();
+    this.extraInheritedEnv = new TreeSet<>();
     this.additionalTools = new ImmutableList.Builder<>();
   }
 
@@ -154,6 +158,11 @@ public TestActionBuilder addExtraEnv(Map<String, String> extraEnv) {
     return this;
   }
 
+  public TestActionBuilder addExtraInheritedEnv(List<String> extraInheritedEnv) {
+    this.extraInheritedEnv.addAll(extraInheritedEnv);
+    return this;
+  }
+
   /**
    * Set the explicit shard count. Note that this may be overridden by the sharding strategy.
    */
@@ -442,7 +451,9 @@ private TestParams createTestAction(int shards)
                 coverageArtifact,
                 coverageDirectory,
                 testProperties,
-                runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
+                runfilesSupport
+                    .getActionEnvironment()
+                    .addVariables(extraTestEnv, extraInheritedEnv),
                 executionSettings,
                 shard,
                 run,