Don't drop the configuration's action environment

See #3320

PiperOrigin-RevId: 220493622

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

@@ -579,7 +579,7 @@ public SkylarkDict<String, String> getSkylarkSubstitutions() {
 
   @Override
   public SkylarkDict<String, String> getEnv() {
-    return SkylarkDict.copyOf(null, env.getFixedEnv());
+    return SkylarkDict.copyOf(null, env.getFixedEnv().toMap());
   }
 
   @Override

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

@@ -16,8 +16,11 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.collect.CollectionUtils;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.util.Fingerprint;
+import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -44,14 +47,88 @@
  */
 @AutoCodec
 public final class ActionEnvironment {
+
+  /** A map of environment variables. */
+  public interface EnvironmentVariables {
+
+    /**
+     * Returns the environment variables as a map.
+     *
+     * <p>WARNING: this allocations additional objects if the underlying implementation is a {@link
+     * CompoundEnvironmentVariables}; use sparingly.
+     */
+    ImmutableMap<String, String> toMap();
+
+    default boolean isEmpty() {
+      return toMap().isEmpty();
+    }
+
+    default int size() {
+      return toMap().size();
+    }
+  }
+
+  /**
+   * An {@link EnvironmentVariables} that combines variables from two different environments without
+   * allocation a new map.
+   */
+  static class CompoundEnvironmentVariables implements EnvironmentVariables {
+    private final EnvironmentVariables current;
+    private final EnvironmentVariables base;
+
+    CompoundEnvironmentVariables(Map<String, String> vars, EnvironmentVariables base) {
+      this.current = new SimpleEnvironmentVariables(vars);
+      this.base = base;
+    }
+
+    @Override
+    public boolean isEmpty() {
+      return current.isEmpty() && base.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);
+    }
+  }
+
+  /** A simple {@link EnvironmentVariables}. */
+  static class SimpleEnvironmentVariables implements EnvironmentVariables {
+
+    static EnvironmentVariables create(Map<String, String> vars) {
+      if (vars.isEmpty()) {
+        return EMPTY_ENVIRONMENT_VARIABLES;
+      }
+      return new SimpleEnvironmentVariables(vars);
+    }
+
+    private final ImmutableMap<String, String> vars;
+
+    private SimpleEnvironmentVariables(Map<String, String> vars) {
+      this.vars = ImmutableMap.copyOf(vars);
+    }
+
+    @Override
+    public ImmutableMap<String, String> toMap() {
+      return vars;
+    }
+  }
+
+  /** An empty {@link EnvironmentVariables}. */
+  public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
+      new SimpleEnvironmentVariables(ImmutableMap.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(ImmutableMap.of(), ImmutableSet.of());
+      new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES, ImmutableSet.of());
 
   /**
    * Splits the given map into a map of variables with a fixed value, and a set of variables that
@@ -71,14 +148,14 @@ public static ActionEnvironment split(Map<String, String> env) {
         inheritedEnv.add(key);
       }
     }
-    return create(ImmutableMap.copyOf(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
+    return create(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
   }
 
-  private final ImmutableMap<String, String> fixedEnv;
-  private final ImmutableSet<String> inheritedEnv;
+  private final EnvironmentVariables fixedEnv;
+  private final Iterable<String> inheritedEnv;
 
-  private ActionEnvironment(
-      ImmutableMap<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
+  private ActionEnvironment(EnvironmentVariables fixedEnv, Iterable<String> inheritedEnv) {
+    CollectionUtils.checkImmutable(inheritedEnv);
     this.fixedEnv = fixedEnv;
     this.inheritedEnv = inheritedEnv;
   }
@@ -90,28 +167,41 @@ private ActionEnvironment(
    */
   @AutoCodec.Instantiator
   public static ActionEnvironment create(
-      ImmutableMap<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
-    if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
+      EnvironmentVariables fixedEnv, Iterable<String> inheritedEnv) {
+    if (fixedEnv.isEmpty() && Iterables.isEmpty(inheritedEnv)) {
       return EMPTY;
     }
     return new ActionEnvironment(fixedEnv, inheritedEnv);
   }
 
+  public static ActionEnvironment create(
+      Map<String, String> fixedEnv, Iterable<String> inheritedEnv) {
+    return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv);
+  }
+
   public static ActionEnvironment create(Map<String, String> fixedEnv) {
-    return new ActionEnvironment(ImmutableMap.copyOf(fixedEnv), ImmutableSet.of());
+    return new ActionEnvironment(new SimpleEnvironmentVariables(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);
   }
 
   /** Returns the combined size of the fixed and inherited environments. */
   public int size() {
-    return fixedEnv.size() + inheritedEnv.size();
+    return fixedEnv.size() + Iterables.size(inheritedEnv);
   }
 
   /**
    * 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 ImmutableMap<String, String> getFixedEnv() {
+  public EnvironmentVariables getFixedEnv() {
     return fixedEnv;
   }
 
@@ -121,7 +211,7 @@ public ImmutableMap<String, String> getFixedEnv() {
    * be used for testing and to compute the cache keys of actions. Use {@link #resolve} instead to
    * get the complete environment.
    */
-  public ImmutableSet<String> getInheritedEnv() {
+  public Iterable<String> getInheritedEnv() {
     return inheritedEnv;
   }
 
@@ -133,7 +223,7 @@ public ImmutableSet<String> getInheritedEnv() {
    */
   public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
     Preconditions.checkNotNull(clientEnv);
-    result.putAll(fixedEnv);
+    result.putAll(fixedEnv.toMap());
     for (String var : inheritedEnv) {
       String value = clientEnv.get(var);
       if (value != null) {
@@ -143,6 +233,7 @@ public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
   }
 
   public void addTo(Fingerprint f) {
-    f.addStringMap(fixedEnv);
+    f.addStringMap(fixedEnv.toMap());
+    f.addStrings(inheritedEnv);
   }
 }

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

@@ -225,7 +225,6 @@ public CommandLines getCommandLines() {
   }
 
   @Override
-  @VisibleForTesting
   public List<String> getArguments() throws CommandLineExpansionException {
     return ImmutableList.copyOf(commandLines.allArguments());
   }
@@ -385,7 +384,7 @@ protected void computeKey(ActionKeyContext actionKeyContext, Fingerprint fp)
     for (Artifact runfilesManifest : runfilesManifests) {
       fp.addPath(runfilesManifest.getExecPath());
     }
-    fp.addStringMap(env.getFixedEnv());
+    env.addTo(fp);
     fp.addStrings(getClientEnvironmentVariables());
     fp.addStringMap(getExecutionInfo());
   }
@@ -395,7 +394,7 @@ public String describeKey() {
     StringBuilder message = new StringBuilder();
     message.append(getProgressMessage());
     message.append('\n');
-    for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
+    for (Map.Entry<String, String> entry : env.getFixedEnv().toMap().entrySet()) {
       message.append("  Environment variable: ");
       message.append(ShellEscaper.escapeString(entry.getKey()));
       message.append('=');
@@ -482,7 +481,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();
+    return env.getFixedEnv().toMap();
   }
 
   /**
@@ -561,6 +560,7 @@ public static class Builder {
     private final List<RunfilesSupplier> inputRunfilesSuppliers = new ArrayList<>();
     private final List<RunfilesSupplier> toolRunfilesSuppliers = new ArrayList<>();
     private ResourceSet resourceSet = AbstractAction.DEFAULT_RESOURCE_SET;
+    private ActionEnvironment actionEnvironment = null;
     private ImmutableMap<String, String> environment = ImmutableMap.of();
     private ImmutableSet<String> inheritedEnvironment = ImmutableSet.of();
     private ImmutableMap<String, String> executionInfo = ImmutableMap.of();
@@ -590,6 +590,7 @@ public Builder(Builder other) {
       this.inputRunfilesSuppliers.addAll(other.inputRunfilesSuppliers);
       this.toolRunfilesSuppliers.addAll(other.toolRunfilesSuppliers);
       this.resourceSet = other.resourceSet;
+      this.actionEnvironment = other.actionEnvironment;
       this.environment = other.environment;
       this.executionInfo = other.executionInfo;
       this.isShellCommand = other.isShellCommand;
@@ -632,9 +633,11 @@ public Action[] build(ActionOwner owner, BuildConfiguration configuration) {
       }
       CommandLines commandLines = result.build();
       ActionEnvironment env =
-          useDefaultShellEnvironment
-              ? configuration.getActionEnvironment()
-              : ActionEnvironment.create(environment, inheritedEnvironment);
+          actionEnvironment != null
+              ? actionEnvironment
+              : useDefaultShellEnvironment
+                  ? configuration.getActionEnvironment()
+                  : ActionEnvironment.create(environment, inheritedEnvironment);
       Action spawnAction =
           buildSpawnAction(
               owner, commandLines, configuration.getCommandLineLimits(), configuration, env);
@@ -825,6 +828,12 @@ public Builder setResources(ResourceSet resourceSet) {
       return this;
     }
 
+    /** Sets the action environment. */
+    public Builder setEnvironment(ActionEnvironment actionEnvironment) {
+      this.actionEnvironment = actionEnvironment;
+      return this;
+    }
+
     /**
      * Sets the map of environment variables. Do not use! This makes the builder ignore the
      * 'default shell environment', which is computed from the --action_env command line option.

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

@@ -99,8 +99,7 @@ protected void computeKey(ActionKeyContext actionKeyContext, Fingerprint fp) {
     fp.addString(GUID);
     fp.addBoolean(filesetTree);
     fp.addBoolean(enableRunfiles);
-    fp.addStringMap(env.getFixedEnv());
-    fp.addStrings(env.getInheritedEnv());
+    env.addTo(fp);
   }
 
   @Override

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

@@ -1558,7 +1558,7 @@ public ActionEnvironment getActionEnvironment() {
    */
   @Override
   public ImmutableMap<String, String> getLocalShellEnvironment() {
-    return actionEnv.getFixedEnv();
+    return actionEnv.getFixedEnv().toMap();
   }
 
   /**
@@ -1576,7 +1576,7 @@ public ImmutableMap<String, String> getLocalShellEnvironment() {
    * client environment.)
    */
   @Deprecated // Use getActionEnvironment instead.
-  public ImmutableSet<String> getVariableShellEnvironment() {
+  public Iterable<String> getVariableShellEnvironment() {
     return actionEnv.getInheritedEnv();
   }
 
@@ -1711,7 +1711,7 @@ public boolean legacyExternalRunfiles() {
    */
   @Override
   public ImmutableMap<String, String> getTestEnv() {
-    return testEnv.getFixedEnv();
+    return testEnv.getFixedEnv().toMap();
   }
 
   /**

src/main/java/com/google/devtools/build/lib/query2/ActionGraphTextOutputFormatterCallback.java

@@ -16,7 +16,7 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
@@ -174,13 +174,13 @@ private void writeAction(ActionAnalysisMetadata action, PrintStream printStream)
       SpawnAction spawnAction = (SpawnAction) action;
       // TODO(twerth): This handles the fixed environment. We probably want to output the inherited
       // environment as well.
-      ImmutableSet<Entry<String, String>> fixedEnvironment =
-          spawnAction.getEnvironment().getFixedEnv().entrySet();
-      if (!fixedEnvironment.isEmpty()) {
+      Iterable<Map.Entry<String, String>> fixedEnvironment =
+          spawnAction.getEnvironment().getFixedEnv().toMap().entrySet();
+      if (!Iterables.isEmpty(fixedEnvironment)) {
         stringBuilder
             .append("  Environment: [")
             .append(
-                fixedEnvironment.stream()
+                Streams.stream(fixedEnvironment)
                     .map(
                         environmentVariable ->
                             environmentVariable.getKey() + "=" + environmentVariable.getValue())

src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java

@@ -1027,8 +1027,7 @@ public boolean isShareable() {
   @Override
   public void computeKey(ActionKeyContext actionKeyContext, Fingerprint fp) {
     fp.addUUID(actionClassId);
-    fp.addStringMap(env.getFixedEnv());
-    fp.addStrings(env.getInheritedEnv());
+    env.addTo(fp);
     fp.addStringMap(compileCommandLine.getEnvironment());
     fp.addStringMap(executionInfo);
 

src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendAction.java

@@ -203,8 +203,7 @@ protected void computeKey(ActionKeyContext actionKeyContext, Fingerprint fp) {
       }
       fp.addPath(imports.getExecPath());
     }
-    fp.addStringMap(env.getFixedEnv());
-    fp.addStrings(env.getInheritedEnv());
+    env.addTo(fp);
     fp.addStringMap(getExecutionInfo());
   }
 

src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileActionBuilder.java

@@ -282,7 +282,8 @@ public void build(RuleContext ruleContext, JavaSemantics javaSemantics) {
     builder.setProgressMessage(getProgressMessage());
     builder.setMnemonic(MNEMONIC);
     builder.setResources(LOCAL_RESOURCES);
-    builder.setEnvironment(UTF8_ENVIRONMENT);
+    builder.setEnvironment(
+        ruleContext.getConfiguration().getActionEnvironment().addFixedVariables(UTF8_ENVIRONMENT));
     builder.setExecutionInfo(executionInfo);
 
     builder.setExtraActionInfo(

src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileActionBuilder.java

@@ -16,6 +16,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.devtools.build.lib.rules.java.JavaCompileActionBuilder.UTF8_ENVIRONMENT;
 import static java.nio.charset.StandardCharsets.ISO_8859_1;
 import static java.util.stream.Collectors.joining;
 
@@ -230,7 +231,8 @@ public void build(JavaToolchainProvider javaToolchain, JavaRuntimeInfo hostJavab
 
     SpawnAction.Builder builder = new SpawnAction.Builder();
 
-    builder.setEnvironment(JavaCompileActionBuilder.UTF8_ENVIRONMENT);
+    builder.setEnvironment(
+        ruleContext.getConfiguration().getActionEnvironment().addFixedVariables(UTF8_ENVIRONMENT));
 
     builder.setProgressMessage(
         new ProgressMessage(