Slightly refactor SpawnAction to improve env handling

This is in preparation for fixing env handling as well as cache key (to use
env) computations in subclasses of SpawnAction.

PiperOrigin-RevId: 196626495

4 files changed, 53 insertions, 21 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java
index 5524e29077..f49ff02094 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java
@@ -36,6 +36,11 @@ import java.util.TreeSet;
  * <p>Inherited environment variables must be declared in the Action interface (see {@link
  * Action#getClientEnvironmentVariables}), so that the dependency on the client environment is known
  * to the execution framework for correct incremental builds.
+ *
+ * <p>By splitting the environment, we can handle environment variable changes more efficiently -
+ * the dependency of the action on the environment variable are tracked in Skyframe (and in the
+ * action cache), such that Bazel knows exactly which actions it needs to rerun, and does not have
+ * to reanalyze the entire dependency graph.
  */
 @AutoCodec
 public final class ActionEnvironment {
@@ -96,10 +101,26 @@ public final class ActionEnvironment {
     return new ActionEnvironment(ImmutableMap.copyOf(fixedEnv), ImmutableSet.of());
   }
 
+  /** Returns the combined size of the fixed and inherited environments. */
+  public int size() {
+    return fixedEnv.size() + inheritedEnv.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 ImmutableMap<String, String> getFixedEnv() {
     return fixedEnv;
   }
 
+  /**
+   * Returns the 'inherited' part of the environment, i.e., those environment variables that are
+   * inherited from the client environment and therefore have no fixed value here. 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 ImmutableSet<String> getInheritedEnv() {
     return inheritedEnv;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index 772ca0e4de..9b826d0906 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -19,6 +19,7 @@ import com.google.common.base.CharMatcher;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.devtools.build.lib.actions.AbstractAction;
@@ -339,9 +340,10 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
    * which also depends on the client environment. Subclasses that which to override the way to get
    * a spawn should override the other GetSpawn() methods instead.
    */
+  @VisibleForTesting
   public final Spawn getSpawn() throws CommandLineExpansionException {
     return new ActionSpawn(
-        commandLines.allArguments(), null, ImmutableList.of(), ImmutableMap.of());
+        commandLines.allArguments(), ImmutableMap.of(), ImmutableList.of(), ImmutableMap.of());
   }
 
   /**
@@ -385,7 +387,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
     for (Artifact runfilesManifest : runfilesManifests) {
       fp.addPath(runfilesManifest.getExecPath());
     }
-    fp.addStringMap(getEnvironment());
+    fp.addStringMap(env.getFixedEnv());
     fp.addStrings(getClientEnvironmentVariables());
     fp.addStringMap(getExecutionInfo());
   }
@@ -395,7 +397,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
     StringBuilder message = new StringBuilder();
     message.append(getProgressMessage());
     message.append('\n');
-    for (Map.Entry<String, String> entry : getEnvironment().entrySet()) {
+    for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
       message.append("  Environment variable: ");
       message.append(ShellEscaper.escapeString(entry.getKey()));
       message.append('=');
@@ -476,7 +478,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
   }
 
   @Override
-  public ImmutableMap<String, String> getEnvironment() {
+  public final ImmutableMap<String, String> getEnvironment() {
     // TODO(ulfjack): AbstractAction should declare getEnvironment with a return value of type
     // ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
     // requires first updating all subclasses and callers to actually handle environments correctly,
@@ -532,17 +534,8 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
       inputs.addAll(additionalInputs);
       this.inputs = inputs.build();
       this.filesetMappings = filesetMappings;
-      LinkedHashMap<String, String> env = new LinkedHashMap<>(SpawnAction.this.getEnvironment());
-      if (clientEnv != null) {
-        for (String var : SpawnAction.this.getClientEnvironmentVariables()) {
-          String value = clientEnv.get(var);
-          if (value == null) {
-            env.remove(var);
-          } else {
-            env.put(var, value);
-          }
-        }
-      }
+      LinkedHashMap<String, String> env = new LinkedHashMap<>(SpawnAction.this.env.size());
+      SpawnAction.this.env.resolve(env, clientEnv);
       effectiveEnvironment = ImmutableMap.copyOf(env);
     }
 
@@ -575,6 +568,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
     private final List<RunfilesSupplier> toolRunfilesSuppliers = new ArrayList<>();
     private ResourceSet resourceSet = AbstractAction.DEFAULT_RESOURCE_SET;
     private ImmutableMap<String, String> environment = ImmutableMap.of();
+    private ImmutableSet<String> inheritedEnvironment = ImmutableSet.of();
     private ImmutableMap<String, String> executionInfo = ImmutableMap.of();
     private boolean isShellCommand = false;
     private boolean useDefaultShellEnvironment = false;
@@ -673,7 +667,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
       ActionEnvironment env =
           useDefaultShellEnvironment
               ? configuration.getActionEnvironment()
-              : ActionEnvironment.create(this.environment);
+              : ActionEnvironment.create(environment, inheritedEnvironment);
       Action spawnAction =
           buildSpawnAction(owner, commandLines, configuration.getCommandLineLimits(), env);
       actions[0] = spawnAction;
@@ -692,7 +686,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
           owner,
           result.build(),
           CommandLineLimits.UNLIMITED,
-          ActionEnvironment.create(this.environment));
+          ActionEnvironment.create(environment, inheritedEnvironment));
     }
 
     private CommandLine buildCommandLinesAndParamFileActions(
@@ -955,6 +949,16 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
     }
 
     /**
+     * Sets the set of inherited environment variables. Do not use! This makes the builder ignore
+     * the 'default shell environment', which is computed from the --action_env command line option.
+     */
+    public Builder setInheritedEnvironment(Iterable<String> inheritedEnvironment) {
+      this.inheritedEnvironment = ImmutableSet.copyOf(inheritedEnvironment);
+      this.useDefaultShellEnvironment = false;
+      return this;
+    }
+
+    /**
      * Sets the map of execution info.
      */
     public Builder setExecutionInfo(Map<String, String> info) {
@@ -995,7 +999,8 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
      */
     public Builder useDefaultShellEnvironment() {
       this.environment = null;
-      this.useDefaultShellEnvironment  = true;
+      this.inheritedEnvironment = null;
+      this.useDefaultShellEnvironment = true;
       return this;
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendAction.java
index 857e457a22..8c5e0561c8 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/LtoBackendAction.java
@@ -203,7 +203,8 @@ public final class LtoBackendAction extends SpawnAction {
       }
       fp.addPath(imports.getExecPath());
     }
-    fp.addStringMap(getEnvironment());
+    fp.addStringMap(env.getFixedEnv());
+    fp.addStrings(env.getInheritedEnv());
     fp.addStringMap(getExecutionInfo());
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java b/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java
index 3e0d629285..25c7e0d444 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java
@@ -34,6 +34,7 @@ import com.google.devtools.build.lib.exec.BinTools;
 import com.google.devtools.build.lib.exec.util.TestExecutorBuilder;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import org.junit.Before;
@@ -157,7 +158,8 @@ public class LtoBackendActionTest extends BuildViewTestCase {
     MNEMONIC,
     RUNFILES_SUPPLIER,
     INPUT,
-    ENVIRONMENT
+    FIXED_ENVIRONMENT,
+    VARIABLE_ENVIRONMENT
   }
 
   @Test
@@ -204,10 +206,13 @@ public class LtoBackendActionTest extends BuildViewTestCase {
             }
 
             Map<String, String> env = new HashMap<>();
-            if (attributesToFlip.contains(KeyAttributes.ENVIRONMENT)) {
+            if (attributesToFlip.contains(KeyAttributes.FIXED_ENVIRONMENT)) {
               env.put("foo", "bar");
             }
             builder.setEnvironment(env);
+            if (attributesToFlip.contains(KeyAttributes.VARIABLE_ENVIRONMENT)) {
+              builder.setInheritedEnvironment(Arrays.asList("baz"));
+            }
 
             Action[] actions =
                 builder.build(