Add ActionEnvironment, which is a wrapper class for a fixed+inherited env

We are currently tracking the action environment (computed from --action_env),
and the test action environment (computed from --test_env, and not including
the action env) as two pairs of fields, and a lot of callers haven't been
updated to handle both parts (TestRunnerAction does, but many 'normal' actions
don't). However, both parts have always both be handled, and moving them into
a single abstraction makes it harder to accidentally miss one part.

We'll subsequently need to update all the action implementations to use the
new abstraction, and also update the Skylark API (or bypass it and always apply
the action environment for all actions, except we don't know which
configuration to use).

PiperOrigin-RevId: 160386181

5 files changed, 150 insertions, 73 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
new file mode 100644
index 0000000000..a70fc1fc56
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java
@@ -0,0 +1,101 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.actions;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.util.Fingerprint;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.TreeSet;
+
+/**
+ * Environment variables for build or test actions.
+ *
+ * <p>The action environment consists of two parts.
+ * <ol>
+ *   <li>All the environment variables with a fixed value, stored in a map.
+ *   <li>All the environment variables inherited from the client environment, stored in a set.
+ * </ol>
+ *
+ * <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.
+ */
+public final class ActionEnvironment {
+  /**
+   * 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()) {
+      if (entry.getValue() != null) {
+        fixedEnv.put(entry.getKey(), entry.getValue());
+      } else {
+        String key = entry.getKey();
+        inheritedEnv.add(key);
+      }
+    }
+    return new ActionEnvironment(fixedEnv, inheritedEnv);
+  }
+
+  private final ImmutableMap<String, String> fixedEnv;
+  private final ImmutableSet<String> inheritedEnv;
+
+  /**
+   * 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.
+   */
+  public ActionEnvironment(Map<String, String> fixedEnv, Set<String> inheritedEnv) {
+    this.fixedEnv = ImmutableMap.copyOf(fixedEnv);
+    this.inheritedEnv = ImmutableSet.copyOf(inheritedEnv);
+  }
+
+  public ImmutableMap<String, String> getFixedEnv() {
+    return fixedEnv;
+  }
+
+  public ImmutableSet<String> getInheritedEnv() {
+    return inheritedEnv;
+  }
+
+  /**
+   * Resolves the action environment and adds the resulting entries to the given {@code result} map,
+   * by looking up any inherited env variables in the given {@code clientEnv}.
+   *
+   * <p>We pass in a map to mutate to avoid creating and merging intermediate maps.
+   */
+  public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
+    Preconditions.checkNotNull(clientEnv);
+    result.putAll(fixedEnv);
+    for (String var : inheritedEnv) {
+      String value = clientEnv.get(var);
+      if (value != null) {
+        result.put(var, value);
+      }
+    }
+  }
+
+  public void addTo(Fingerprint f) {
+    f.addStringMap(fixedEnv);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
index 993394a2d3..c9a933b8ac 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
@@ -34,6 +34,7 @@ import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.MutableClassToInstanceMap;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.Root;
 import com.google.devtools.build.lib.analysis.AspectCollection;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
@@ -70,7 +71,6 @@ import com.google.devtools.build.lib.skylarkinterface.SkylarkModule;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkModuleCategory;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.OS;
-import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.util.RegexFilter;
 import com.google.devtools.build.lib.vfs.Path;
@@ -99,7 +99,6 @@ import java.util.Objects;
 import java.util.Queue;
 import java.util.Set;
 import java.util.TreeMap;
-import java.util.TreeSet;
 import javax.annotation.Nullable;
 
 /**
@@ -1306,22 +1305,8 @@ public final class BuildConfiguration implements BuildEvent {
    */
   private final ImmutableMap<String, String> globalMakeEnv;
 
-  // The action and test environments each consist of two parts. The first part are all the
-  // environment variables with a fixed value, and is stored in a map. The second part are all the
-  // environment variables inherited from the client environment, and is stored in a set. The second
-  // part is resolved as part of action execution - they are declared in the Action interface, and
-  // the SkyframeExecutor adds dependency edges from the actions to nodes corresponding to the
-  // individual client env variables, which ensures correct invalidation when the client env
-  // changes.
-  // Care needs to be taken that the two sets don't overlap - the order in which the two parts are
-  // combined later is undefined.
-  // TODO(ulfjack): Consider adding a wrapper class to construct and contain the two parts, and also
-  // handle case independence on certain platforms.
-  private final ImmutableMap<String, String> actionEnv;
-  private final ImmutableSet<String> inheritedActionEnvVars;
-
-  private final ImmutableMap<String, String> testEnv;
-  private final ImmutableSet<String> inheritedTestEnvVars;
+  private final ActionEnvironment actionEnv;
+  private final ActionEnvironment testEnv;
 
   private final BuildOptions buildOptions;
   private final Options options;
@@ -1471,7 +1456,7 @@ public final class BuildConfiguration implements BuildEvent {
    * statically set environment variables with their values and the set of environment variables to
    * be inherited from the client environment.
    */
-  private Pair<ImmutableMap<String, String>, ImmutableSet<String>> setupActionEnvironment() {
+  private ActionEnvironment setupActionEnvironment() {
     // We make a copy first to remove duplicate entries; last one wins.
     Map<String, String> actionEnv = new HashMap<>();
     // TODO(ulfjack): Remove all env variables from configuration fragments.
@@ -1484,7 +1469,7 @@ public final class BuildConfiguration implements BuildEvent {
     for (Map.Entry<String, String> entry : options.actionEnvironment) {
       actionEnv.put(entry.getKey(), entry.getValue());
     }
-    return split(actionEnv);
+    return ActionEnvironment.split(actionEnv);
   }
 
   /**
@@ -1492,27 +1477,13 @@ public final class BuildConfiguration implements BuildEvent {
    * statically set environment variables with their values and the set of environment variables to
    * be inherited from the client environment.
    */
-  private Pair<ImmutableMap<String, String>, ImmutableSet<String>> setupTestEnvironment() {
+  private ActionEnvironment setupTestEnvironment() {
     // We make a copy first to remove duplicate entries; last one wins.
     Map<String, String> testEnv = new HashMap<>();
     for (Map.Entry<String, String> entry : options.testEnvironment) {
       testEnv.put(entry.getKey(), entry.getValue());
     }
-    return split(testEnv);
-  }
-
-  private Pair<ImmutableMap<String, String>, ImmutableSet<String>> split(Map<String, String> env) {
-    Map<String, String> fixedEnv = new TreeMap<>();
-    Set<String> variableEnv = new TreeSet<>();
-    for (Map.Entry<String, String> entry : env.entrySet()) {
-      if (entry.getValue() != null) {
-        fixedEnv.put(entry.getKey(), entry.getValue());
-      } else {
-        String key = entry.getKey();
-        variableEnv.add(key);
-      }
-    }
-    return Pair.of(ImmutableMap.copyOf(fixedEnv), ImmutableSet.copyOf(variableEnv));
+    return ActionEnvironment.split(testEnv);
   }
 
   /**
@@ -1583,13 +1554,9 @@ public final class BuildConfiguration implements BuildEvent {
 
     this.shellExecutable = computeShellExecutable();
 
-    Pair<ImmutableMap<String, String>, ImmutableSet<String>> actionEnvs = setupActionEnvironment();
-    this.actionEnv = actionEnvs.getFirst();
-    this.inheritedActionEnvVars = actionEnvs.getSecond();
+    this.actionEnv = setupActionEnvironment();
 
-    Pair<ImmutableMap<String, String>, ImmutableSet<String>> testEnvs = setupTestEnvironment();
-    this.testEnv = testEnvs.getFirst();
-    this.inheritedTestEnvVars = testEnvs.getSecond();
+    this.testEnv = setupTestEnvironment();
 
     this.transitiveOptionDetails = computeOptionsMap(buildOptions, fragments.values());
 
@@ -2353,6 +2320,10 @@ public final class BuildConfiguration implements BuildEvent {
     return checksum();
   }
 
+  public ActionEnvironment getActionEnvironment() {
+    return actionEnv;
+  }
+
   @SkylarkCallable(
     name = "default_shell_env",
     structField = true,
@@ -2371,8 +2342,9 @@ public final class BuildConfiguration implements BuildEvent {
    * <p>Since values of the "fixed" variables are already known at analysis phase, it is returned
    * here as a map.
    */
+  @Deprecated // Use getActionEnvironment instead.
   public ImmutableMap<String, String> getLocalShellEnvironment() {
-    return actionEnv;
+    return actionEnv.getFixedEnv();
   }
 
   /**
@@ -2389,8 +2361,9 @@ public final class BuildConfiguration implements BuildEvent {
    * environment. (Variables where the name is not returned in this set should not be taken from the
    * client environment.)
    */
+  @Deprecated // Use getActionEnvironment instead.
   public ImmutableSet<String> getVariableShellEnvironment() {
-    return inheritedActionEnvVars;
+    return actionEnv.getInheritedEnv();
   }
 
   /**
@@ -2540,19 +2513,27 @@ public final class BuildConfiguration implements BuildEvent {
    * Returns user-specified test environment variables and their values, as set by the --test_env
    * options.
    */
+  @Deprecated
   @SkylarkCallable(
     name = "test_env",
     structField = true,
     doc =
         "A dictionary containing user-specified test environment variables and their values, "
-            + "as set by the --test_env options."
+            + "as set by the --test_env options. DO NOT USE! This is not the complete environment!"
   )
   public ImmutableMap<String, String> getTestEnv() {
-    return testEnv;
+    return testEnv.getFixedEnv();
   }
 
-  public ImmutableSet<String> getInheritedTestEnv() {
-    return inheritedTestEnvVars;
+  /**
+   * Returns user-specified test environment variables and their values, as set by the
+   * {@code --test_env} options. It is incomplete in that it is not a superset of the
+   * {@link #getActionEnvironment}, but both have to be applied, with this one being applied after
+   * the other, such that {@code --test_env} settings can override {@code --action_env} settings.
+   */
+  // TODO(ulfjack): Just return the merged action and test action environment here?
+  public ActionEnvironment getTestActionEnvironment() {
+    return testEnv;
   }
 
   public TriState cacheTestResults() {
diff --git a/src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java b/src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java
index 5685601008..09bf0e04b3 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java
@@ -88,22 +88,10 @@ public class TestPolicy {
     env.putAll(testAction.getExtraTestEnv());
 
     // Overwrite with the environment common to all actions, see --action_env.
-    env.putAll(testAction.getConfiguration().getLocalShellEnvironment());
-    for (String key : testAction.getConfiguration().getVariableShellEnvironment()) {
-      String value = clientEnv.get(key);
-      if (value != null) {
-        env.put(key, value);
-      }
-    }
+    testAction.getConfiguration().getActionEnvironment().resolve(env, clientEnv);
 
     // Overwrite with the environment common to all tests, see --test_env.
-    env.putAll(testAction.getConfiguration().getTestEnv());
-    for (String key : testAction.getConfiguration().getInheritedTestEnv()) {
-      String value = clientEnv.get(key);
-      if (value != null) {
-        env.put(key, value);
-      }
-    }
+    testAction.getConfiguration().getTestActionEnvironment().resolve(env, clientEnv);
 
     // Setup any test-specific env variables; note that this does not overwrite existing values for
     // TEST_RANDOM_SEED or TEST_SIZE if they're already set.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/TestSupport.java b/src/main/java/com/google/devtools/build/lib/rules/objc/TestSupport.java
index a416d7630f..b11c77eadb 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/TestSupport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/TestSupport.java
@@ -76,7 +76,10 @@ public class TestSupport {
     String runMemleaks =
         ruleContext.getFragment(ObjcConfiguration.class).runMemleaks() ? "true" : "false";
 
-    ImmutableMap<String, String> testEnv = ruleContext.getConfiguration().getTestEnv();
+    // TODO(ulfjack): This is missing the action environment, and the inherited parts from both. Is
+    // that intentional? We should either fix it, or clearly document why we're doing that.
+    ImmutableMap<String, String> testEnv =
+        ruleContext.getConfiguration().getTestActionEnvironment().getFixedEnv();
 
     // The substitutions below are common for simulator and lab device.
     ImmutableList.Builder<Substitution> substitutions =
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/TestRunnerAction.java b/src/main/java/com/google/devtools/build/lib/rules/test/TestRunnerAction.java
index a9349db769..7726361a49 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/TestRunnerAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/TestRunnerAction.java
@@ -32,6 +32,7 @@ import com.google.devtools.build.lib.analysis.RunfilesSupplierImpl;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.RunUnder;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.collect.ImmutableIterable;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.LoggingUtil;
@@ -92,15 +93,17 @@ public class TestRunnerAction extends AbstractAction implements NotifyOnActionCa
   private final String workspaceName;
   private final boolean useTestRunner;
 
-  // Mutable state related to test caching.
-  private Boolean unconditionalExecution; // lazily initialized: null indicates unknown
+  // Mutable state related to test caching. Lazily initialized: null indicates unknown.
+  private Boolean unconditionalExecution;
 
-  // Any extra environment variables (and values) added by the rule that created this action.
+  /** Any extra environment variables (and values) added by the rule that created this action. */
   private final ImmutableMap<String, String> extraTestEnv;
 
-  // These are handled explicitly by the ActionCacheChecker and so don't have to be included in the
-  // cache key.
-  private final Iterable<String> requiredClientEnvVariables;
+  /**
+   * The set of environment variables that are inherited from the client environment. These are
+   * handled explicitly by the ActionCacheChecker and so don't have to be included in the cache key.
+   */
+  private final ImmutableIterable<String> requiredClientEnvVariables;
 
   private static ImmutableList<Artifact> list(Artifact... artifacts) {
     ImmutableList.Builder<Artifact> builder = ImmutableList.builder();
@@ -179,8 +182,9 @@ public class TestRunnerAction extends AbstractAction implements NotifyOnActionCa
 
     this.extraTestEnv = ImmutableMap.copyOf(extraTestEnv);
     this.requiredClientEnvVariables =
-        Iterables.concat(
-            configuration.getVariableShellEnvironment(), configuration.getInheritedTestEnv());
+        ImmutableIterable.from(Iterables.concat(
+            configuration.getActionEnvironment().getInheritedEnv(),
+            configuration.getTestActionEnvironment().getInheritedEnv()));
   }
 
   public BuildConfiguration getConfiguration() {
@@ -227,8 +231,8 @@ public class TestRunnerAction extends AbstractAction implements NotifyOnActionCa
     f.addStringMap(extraTestEnv);
     // TODO(ulfjack): It might be better for performance to hash the action and test envs in config,
     // and only add a hash here.
-    f.addStringMap(configuration.getLocalShellEnvironment());
-    f.addStringMap(configuration.getTestEnv());
+    configuration.getActionEnvironment().addTo(f);
+    configuration.getTestActionEnvironment().addTo(f);
     // The 'requiredClientEnvVariables' are handled by Skyframe and don't need to be added here.
     f.addString(testProperties.getSize().toString());
     f.addString(testProperties.getTimeout().toString());