Let Starlark tests inherit env variables

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

@@ -699,7 +699,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

@@ -19,6 +19,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.util.Fingerprint;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -45,36 +46,57 @@
  */
 public final class ActionEnvironment {
 
-  /** A map of environment variables. */
+  /**
+   * A map of environment variables and their values together with a list of variables whose values
+   * should be inherited from the client environment.
+   */
   public interface EnvironmentVariables {
 
     /**
-     * Returns the environment variables as a map.
+     * Returns the fixed environment variables as a map.
+     *
+     * <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
+     * CompoundEnvironmentVariables}; use sparingly.
+     */
+    ImmutableMap<String, String> getFixedEnvironment();
+
+    /**
+     * Returns the inherited environment variables as a set.
      *
-     * <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();
+    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();
     }
   }
 
   /**
    * An {@link EnvironmentVariables} that combines variables from two different environments without
-   * allocation a new map.
+   * allocating 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()) {
+        return base;
+      }
+      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,57 +106,73 @@ public boolean isEmpty() {
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
+    public ImmutableMap<String, String> getFixedEnvironment() {
       Map<String, String> result = new LinkedHashMap<>();
-      result.putAll(base.toMap());
-      result.putAll(current.toMap());
+      result.putAll(base.getFixedEnvironment());
+      result.putAll(current.getFixedEnvironment());
       return ImmutableMap.copyOf(result);
     }
+
+    @Override
+    public ImmutableSet<String> getInheritedEnvironment() {
+      Set<String> result = new LinkedHashSet<>();
+      result.addAll(base.getInheritedEnvironment());
+      result.addAll(current.getInheritedEnvironment());
+      return ImmutableSet.copyOf(result);
+    }
   }
 
-  /** A simple {@link EnvironmentVariables}. */
+  /**
+   * 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> fixedVars, Set<String> inheritedVars) {
+      this.fixedVars = ImmutableMap.copyOf(fixedVars);
+      this.inheritedVars = ImmutableSet.copyOf(inheritedVars);
+    }
 
-    private SimpleEnvironmentVariables(Map<String, String> vars) {
-      this.vars = ImmutableMap.copyOf(vars);
+    @Override
+    public ImmutableMap<String, String> getFixedEnvironment() {
+      return fixedVars;
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
-      return vars;
+    public ImmutableSet<String> getInheritedEnvironment() {
+      return inheritedVars;
     }
   }
 
-  /** An empty {@link EnvironmentVariables}. */
+  /**
+   * 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
    * should be inherited, the latter of which are identified by having a {@code null} value in the
    * given map. Returns these two parts as a new {@link ActionEnvironment} instance.
    */
   public static ActionEnvironment split(Map<String, String> env) {
-    // Care needs to be taken that the two sets don't overlap - the order in which the two parts are
-    // combined later is undefined.
     Map<String, String> fixedEnv = new TreeMap<>();
     Set<String> inheritedEnv = new TreeSet<>();
     for (Map.Entry<String, String> entry : env.entrySet()) {
@@ -145,59 +183,81 @@ 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.
+   * Creates a new action environment. If an environment variable is contained in both {@link
+   * EnvironmentVariables#getFixedEnvironment()} and {@link EnvironmentVariables#getInheritedEnvironment()},
+   * the result of {@link ActionEnvironment#resolve(Map, Map)} will contain the value inherited from
+   * the client environment.
    */
-  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);
   }
 
+  /**
+   * Creates a new action environment. If an environment variable is contained both as a key in
+   * {@code fixedEnv} and in {@code inheritedEnv}, the result of {@link
+   * ActionEnvironment#resolve(Map, Map)} will contain the value inherited from the client
+   * environment.
+   */
   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 withAdditionalFixedVariables(Map<String, String> fixedVars) {
+    return withAdditionalVariables(fixedVars, ImmutableSet.of());
   }
 
-  /** Returns the combined size of the fixed and inherited environments. */
+  /**
+   * 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 withAdditionalVariables(Map<String, String> fixedVars,
+      Set<String> inheritedVars) {
+    EnvironmentVariables newVars = CompoundEnvironmentVariables.create(fixedVars, inheritedVars,
+        vars);
+    if (newVars == vars) {
+      return this;
+    }
+    return ActionEnvironment.create(newVars);
+  }
+
+  /**
+   * Returns an upper bound on the combined size of the fixed and inherited environments. A call to
+   * {@link ActionEnvironment#resolve(Map, Map)} may add less entries than this number if
+   * environment variables are contained in both the fixed and the inherited environment.
+   */
   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 +267,7 @@ public EnvironmentVariables getFixedEnv() {
    * get the complete environment.
    */
   public ImmutableSet<String> getInheritedEnv() {
-    return inheritedEnv;
+    return vars.getInheritedEnvironment();
   }
 
   /**
@@ -218,8 +278,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 +288,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

@@ -475,6 +475,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

@@ -465,7 +465,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('=');
@@ -574,7 +574,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();
   }
 
   /** Returns the out-of-band execution data for this action. */

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

@@ -485,7 +485,7 @@ public boolean isSiblingRepositoryLayout() {
    */
   @Override
   public ImmutableMap<String, String> getLocalShellEnvironment() {
-    return actionEnv.getFixedEnv().toMap();
+    return actionEnv.getFixedEnv();
   }
 
   /**
@@ -629,7 +629,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

@@ -54,7 +54,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;
 
 /**
@@ -94,10 +96,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<>();
   }
 
@@ -147,6 +151,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.
    */
@@ -402,7 +411,7 @@ private TestParams createTestAction(int shards)
                 coverageArtifact,
                 coverageDirectory,
                 testProperties,
-                runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
+                runfilesSupport.getActionEnvironment().withAdditionalVariables(extraTestEnv, extraInheritedEnv),
                 executionSettings,
                 shard,
                 run,