Add an env attribute to all test and binary rule classes

Fixes #7364.

PiperOrigin-RevId: 345355968

Author: Googler

src/main/java/com/google/devtools/build/docgen/PredefinedAttributes.java

@@ -27,16 +27,19 @@
 public class PredefinedAttributes {
 
   /**
-   * List of documentation for common attributes of *_test rules, relative to
-   * {@link com.google.devtools.build.docgen}.
+   * List of documentation for common attributes of *_test rules, relative to {@link
+   * com.google.devtools.build.docgen}.
    */
-  public static final ImmutableList<String> TEST_ATTRIBUTES_DOCFILES = ImmutableList.of(
-      "templates/attributes/test/args.html",
-      "templates/attributes/test/size.html",
-      "templates/attributes/test/timeout.html",
-      "templates/attributes/test/flaky.html",
-      "templates/attributes/test/shard_count.html",
-      "templates/attributes/test/local.html");
+  public static final ImmutableList<String> TEST_ATTRIBUTES_DOCFILES =
+      ImmutableList.of(
+          "templates/attributes/test/args.html",
+          "templates/attributes/test/env.html",
+          "templates/attributes/test/env_inherit.html",
+          "templates/attributes/test/size.html",
+          "templates/attributes/test/timeout.html",
+          "templates/attributes/test/flaky.html",
+          "templates/attributes/test/shard_count.html",
+          "templates/attributes/test/local.html");
 
   /**
    * List of common attributes documentation, relative to {@link com.google.devtools.build.docgen}.
@@ -60,12 +63,13 @@ public class PredefinedAttributes {
           "templates/attributes/common/visibility.html");
 
   /**
-   * List of documentation for common attributes of *_binary rules, relative to
-   * {@link com.google.devtools.build.docgen}.
+   * List of documentation for common attributes of *_binary rules, relative to {@link
+   * com.google.devtools.build.docgen}.
    */
   public static final ImmutableList<String> BINARY_ATTRIBUTES_DOCFILES =
       ImmutableList.of(
           "templates/attributes/binary/args.html",
+          "templates/attributes/binary/env.html",
           "templates/attributes/binary/output_licenses.html");
 
   private static ImmutableMap<String, RuleDocumentationAttribute> generateAttributeMap(

src/main/java/com/google/devtools/build/docgen/templates/attributes/binary/env.html

@@ -0,0 +1,13 @@
+<p><code>Dictionary of strings; optional; values are subject to
+<a href="${link make-variables#location}">$(location)</a> and
+<a href="${link make-variables}">"Make variable"</a> substitution</code></p>
+
+<p>Specifies additional environment variables to set when the target is
+  executed by <code>bazel run</code>.
+</p>
+
+<p>
+<em class="harmful">NOTE: The environment variables are not set when you run the target
+outside of bazel (for example, by manually executing the binary in
+<code>bazel-bin/</code>).</em>
+</p>

src/main/java/com/google/devtools/build/docgen/templates/attributes/test/env.html

@@ -0,0 +1,8 @@
+<p><code>Dictionary of strings; optional; values are subject to
+<a href="${link make-variables#location}">$(location)</a> and
+<a href="${link make-variables}">"Make variable"</a> substitution</code>
+</p>
+
+<p>Specifies additional environment variables to set when the test is
+  executed by <code>bazel test</code>.
+</p>

src/main/java/com/google/devtools/build/docgen/templates/attributes/test/env_inherit.html

@@ -0,0 +1,5 @@
+<p><code>List of strings; optional</p>
+
+<p>Specifies additional environment variables to inherit from the
+  external environment when the test is executed by <code>bazel test</code>.
+</p>

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

@@ -194,6 +194,8 @@ public Object getDefault(AttributeMap rule) {
                   .taggable()
                   .nonconfigurable("policy decision: should be consistent across configurations"))
           .add(attr("args", STRING_LIST))
+          .add(attr("env", STRING_DICT))
+          .add(attr("env_inherit", STRING_LIST))
           // Input files for every test action
           .add(
               attr("$test_wrapper", LABEL)
@@ -475,6 +477,7 @@ public static final class BinaryBaseRule implements RuleDefinition {
     public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
       return builder
           .add(attr("args", STRING_LIST))
+          .add(attr("env", STRING_DICT))
           .add(attr("output_licenses", LICENSE))
           .add(
               attr("$is_executable", BOOLEAN)

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

@@ -16,6 +16,8 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLine;
 import com.google.devtools.build.lib.analysis.SourceManifestAction.ManifestType;
@@ -34,9 +36,11 @@
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Collection;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import javax.annotation.Nullable;
 
 /**
@@ -85,6 +89,7 @@ public final class RunfilesSupport {
   private final boolean buildRunfileLinks;
   private final boolean runfilesEnabled;
   private final CommandLine args;
+  private final ActionEnvironment actionEnvironment;
 
   /**
    * Creates the RunfilesSupport helper with the given executable and runfiles.
@@ -94,7 +99,11 @@ public final class RunfilesSupport {
    * @param runfiles the runfiles
    */
   private static RunfilesSupport create(
-      RuleContext ruleContext, Artifact executable, Runfiles runfiles, CommandLine args) {
+      RuleContext ruleContext,
+      Artifact executable,
+      Runfiles runfiles,
+      CommandLine args,
+      ActionEnvironment actionEnvironment) {
     Artifact owningExecutable = Preconditions.checkNotNull(executable);
     boolean createManifest = ruleContext.getConfiguration().buildRunfilesManifests();
     boolean buildRunfileLinks = ruleContext.getConfiguration().buildRunfileLinks();
@@ -139,7 +148,8 @@ private static RunfilesSupport create(
         owningExecutable,
         buildRunfileLinks,
         runfilesEnabled,
-        args);
+        args,
+        actionEnvironment);
   }
 
   @AutoCodec.Instantiator
@@ -152,7 +162,8 @@ private static RunfilesSupport create(
       Artifact owningExecutable,
       boolean buildRunfileLinks,
       boolean runfilesEnabled,
-      CommandLine args) {
+      CommandLine args,
+      ActionEnvironment actionEnvironment) {
     this.runfiles = runfiles;
     this.runfilesInputManifest = runfilesInputManifest;
     this.runfilesManifest = runfilesManifest;
@@ -161,6 +172,7 @@ private static RunfilesSupport create(
     this.buildRunfileLinks = buildRunfileLinks;
     this.runfilesEnabled = runfilesEnabled;
     this.args = args;
+    this.actionEnvironment = actionEnvironment;
   }
 
   /** Returns the executable owning this RunfilesSupport. Only use from Starlark. */
@@ -394,14 +406,23 @@ public CommandLine getArgs() {
     return args;
   }
 
+  /** Returns the immutable environment from the 'env' and 'env_inherit' attribute values. */
+  public ActionEnvironment getActionEnvironment() {
+    return actionEnvironment;
+  }
+
   /**
    * Creates and returns a {@link RunfilesSupport} object for the given rule and executable. Note
    * that this method calls back into the passed in rule to obtain the runfiles.
    */
   public static RunfilesSupport withExecutable(
       RuleContext ruleContext, Runfiles runfiles, Artifact executable) {
     return RunfilesSupport.create(
-        ruleContext, executable, runfiles, computeArgs(ruleContext, CommandLine.EMPTY));
+        ruleContext,
+        executable,
+        runfiles,
+        computeArgs(ruleContext, CommandLine.EMPTY),
+        computeActionEnvironment(ruleContext));
   }
 
   /**
@@ -411,7 +432,11 @@ public static RunfilesSupport withExecutable(
   public static RunfilesSupport withExecutable(
       RuleContext ruleContext, Runfiles runfiles, Artifact executable, List<String> appendingArgs) {
     return RunfilesSupport.create(
-        ruleContext, executable, runfiles, computeArgs(ruleContext, CommandLine.of(appendingArgs)));
+        ruleContext,
+        executable,
+        runfiles,
+        computeArgs(ruleContext, CommandLine.of(appendingArgs)),
+        computeActionEnvironment(ruleContext));
   }
 
   /**
@@ -421,7 +446,11 @@ public static RunfilesSupport withExecutable(
   public static RunfilesSupport withExecutable(
       RuleContext ruleContext, Runfiles runfiles, Artifact executable, CommandLine appendingArgs) {
     return RunfilesSupport.create(
-        ruleContext, executable, runfiles, computeArgs(ruleContext, appendingArgs));
+        ruleContext,
+        executable,
+        runfiles,
+        computeArgs(ruleContext, appendingArgs),
+        computeActionEnvironment(ruleContext));
   }
 
   private static CommandLine computeArgs(RuleContext ruleContext, CommandLine additionalArgs) {
@@ -434,6 +463,30 @@ private static CommandLine computeArgs(RuleContext ruleContext, CommandLine addi
         ruleContext.getExpander().withDataLocations().tokenized("args"), additionalArgs);
   }
 
+  private static ActionEnvironment computeActionEnvironment(RuleContext ruleContext) {
+    if (!ruleContext.getRule().isAttrDefined("env", Type.STRING_DICT)
+        && !ruleContext.getRule().isAttrDefined("env_inherit", Type.STRING_LIST)) {
+      return ActionEnvironment.EMPTY;
+    }
+    TreeMap<String, String> fixedEnv = new TreeMap<>();
+    Set<String> inheritedEnv = new LinkedHashSet<>();
+    if (ruleContext.isAttrDefined("env", Type.STRING_DICT)) {
+      Expander expander = ruleContext.getExpander().withDataLocations();
+      for (Map.Entry<String, String> entry :
+          ruleContext.attributes().get("env", Type.STRING_DICT).entrySet()) {
+        fixedEnv.put(entry.getKey(), expander.expand("env", entry.getValue()));
+      }
+    }
+    if (ruleContext.isAttrDefined("env_inherit", Type.STRING_LIST)) {
+      for (String key : ruleContext.attributes().get("env_inherit", Type.STRING_LIST)) {
+        if (!fixedEnv.containsKey(key)) {
+          inheritedEnv.add(key);
+        }
+      }
+    }
+    return ActionEnvironment.create(fixedEnv, ImmutableSet.copyOf(inheritedEnv));
+  }
+
   /** Returns the path of the input manifest of {@code runfilesDir}. */
   public static Path inputManifestPath(Path runfilesDir) {
     return FileSystemUtils.replaceExtension(runfilesDir, INPUT_MANIFEST_EXT);

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

@@ -52,7 +52,6 @@
  * Helper class to create test actions.
  */
 public final class TestActionBuilder {
-
   private static final String CC_CODE_COVERAGE_SCRIPT = "CC_CODE_COVERAGE_SCRIPT";
   private static final String LCOV_MERGER = "LCOV_MERGER";
   // The coverage tool Bazel uses to generate a code coverage report for C++.
@@ -396,7 +395,7 @@ private TestParams createTestAction(int shards) throws InterruptedException {
                 coverageArtifact,
                 coverageDirectory,
                 testProperties,
-                extraTestEnv,
+                runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
                 executionSettings,
                 shard,
                 run,

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

@@ -28,6 +28,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.actions.AbstractAction;
 import com.google.devtools.build.lib.actions.ActionContinuationOrResult;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionInput;
@@ -132,7 +133,7 @@ public class TestRunnerAction extends AbstractAction
   private Boolean unconditionalExecution;
 
   /** Any extra environment variables (and values) added by the rule that created this action. */
-  private final ImmutableMap<String, String> extraTestEnv;
+  private final ActionEnvironment extraTestEnv;
 
   /**
    * The set of environment variables that are inherited from the client environment. These are
@@ -176,7 +177,7 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
       Artifact coverageArtifact,
       @Nullable Artifact coverageDirectory,
       TestTargetProperties testProperties,
-      Map<String, String> extraTestEnv,
+      ActionEnvironment extraTestEnv,
       TestTargetExecutionSettings executionSettings,
       int shardNum,
       int runNumber,
@@ -236,12 +237,13 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
     this.testInfrastructureFailure = baseDir.getChild("test.infrastructure_failure");
     this.workspaceName = workspaceName;
 
-    this.extraTestEnv = ImmutableMap.copyOf(extraTestEnv);
+    this.extraTestEnv = extraTestEnv;
     this.requiredClientEnvVariables =
         ImmutableIterable.from(
             Iterables.concat(
                 configuration.getActionEnvironment().getInheritedEnv(),
-                configuration.getTestActionEnvironment().getInheritedEnv()));
+                configuration.getTestActionEnvironment().getInheritedEnv(),
+                this.extraTestEnv.getInheritedEnv()));
     this.cancelConcurrentTestsOnSuccess = cancelConcurrentTestsOnSuccess;
     this.splitCoveragePostProcessing = splitCoveragePostProcessing;
     this.lcovMergerFilesToRun = lcovMergerFilesToRun;
@@ -378,7 +380,7 @@ protected void computeKey(
     fp.addBoolean(executionSettings.getTestRunnerFailFast());
     RunUnder runUnder = executionSettings.getRunUnder();
     fp.addString(runUnder == null ? "" : runUnder.getValue());
-    fp.addStringMap(extraTestEnv);
+    extraTestEnv.addTo(fp);
     // TODO(ulfjack): It might be better for performance to hash the action and test envs in config,
     // and only add a hash here.
     configuration.getActionEnvironment().addTo(fp);
@@ -675,7 +677,7 @@ public Artifact getTestLog() {
   }
 
   /** Returns all environment variables which must be set in order to run this test. */
-  public Map<String, String> getExtraTestEnv() {
+  public ActionEnvironment getExtraTestEnv() {
     return extraTestEnv;
   }
 

src/main/java/com/google/devtools/build/lib/exec/TestPolicy.java

@@ -88,7 +88,7 @@ public Map<String, String> computeTestEnvironment(
     }
 
     // Rule-specified test env.
-    env.putAll(testAction.getExtraTestEnv());
+    testAction.getExtraTestEnv().resolve(env, clientEnv);
 
     // Overwrite with the environment common to all actions, see --action_env.
     testAction.getConfiguration().getActionEnvironment().resolve(env, clientEnv);

src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java

@@ -471,6 +471,9 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
       }
     } else {
       workingDir = runfilesDir;
+      if (runfilesSupport != null) {
+        runfilesSupport.getActionEnvironment().resolve(runEnvironment, env.getClientEnv());
+      }
       List<String> args = computeArgs(env, targetToRun, commandLineArgs);
       try {
         constructCommandLine(

src/test/shell/integration/run_test.sh

@@ -487,6 +487,37 @@ EOF
   assert_starts_with "${tmpdir}/test_bazel_run_with_custom_tmpdir" "$(cat "${TEST_TMPDIR}/tmpdir_value")"
 }
 
+function test_run_binary_with_env_attribute() {
+  local -r pkg="pkg${LINENO}"
+  mkdir -p ${pkg}
+  cat > $pkg/BUILD <<'EOF'
+sh_binary(
+  name = 't',
+  srcs = [':t.sh'],
+  data = [':t.dat'],
+  env = {
+    "ENV_A": "not_inherited",
+    "ENV_C": "no_surprise",
+    "ENV_DATA": "$(location :t.dat)",
+  },
+)
+EOF
+  cat > $pkg/t.sh <<'EOF'
+#!/bin/sh
+env
+cat $ENV_DATA
+exit 0
+EOF
+  touch $pkg/t.dat
+  chmod +x $pkg/t.sh
+  ENV_B=surprise ENV_C=surprise bazel run //$pkg:t > $TEST_log \
+      || fail "expected test to pass"
+  expect_log "ENV_A=not_inherited"
+  expect_log "ENV_B=surprise"
+  expect_log "ENV_C=no_surprise"
+  expect_log "ENV_DATA=$pkg/t.dat"
+}
+
 # Usage: assert_starts_with PREFIX STRING_TO_CHECK.
 # Asserts that `$1` is a prefix of `$2`.
 function assert_starts_with() {