Let Starlark executable rules specify their environment (bazelbuild#15766)

* Specify fixedEnv/inheritedEnv interaction in ActionEnvironment

Previously, ActionEnvironment did not publicly document how fixed and
inherited environment variables interact, but still cautioned users to
keep the two sets disjoint without enforcing this. As a result, neither
could users rely on the interaction nor could ActionEnvironment benefit
from the additional freedom of not specifying the behavior.

With this commit, ActionEnvironment explicitly specifies that the values
of environment variable inherited from the client environment take
precedence over fixed values and codifies this behavior in a test.
This has been the effective behavior all along and has the advantage
that users can provide overrideable defaults for environment variables.

Closes bazelbuild#15170.

PiperOrigin-RevId: 439315634

* Intern trivial ActionEnvironment and EnvironmentVariables instances

When an ActionEnvironment is constructed out of an existing one, the
ActionEnvironment and EnvironmentVariables instances, which are
immutable, can be reused if no variables are added.

Also renames addVariables and addFixedVariables to better reflect that
they are operating on an immutable type.

Closes bazelbuild#15171.

PiperOrigin-RevId: 440312159

* Let Starlark executable rules specify their environment

The new RunEnvironmentInfo provider allows any executable or test
Starlark rule to specify the environment for when it is executed, either
as part of a test action or via the run command.

Refactors testing.TestEnvironment to construct a RunEnvironmentInfo and
adds a warning (but not an error) if the provider constructed in this
way is returned from a non-executable non-test rule. If a
RunEnvironmentInfo is constructed directly via the Starlark constructor,
this warning becomes an error.

Fixes bazelbuild#7364
Fixes bazelbuild#15224
Fixes bazelbuild#15225

Closes bazelbuild#15232.

PiperOrigin-RevId: 448185352

* Fix strict deps violation

```
src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java:990: error: [strict] Using type com.google.devtools.build.lib.cmdline.LabelValidator from an indirect dependency (TOOL_INFO: "//src/main/java/com/google/devtools/build/lib/cmdline:LabelValidator"). See command below **
      LabelValidator.parseAbsoluteLabel(labelString);
      ^
```

Author: fmeum

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

@@ -46,8 +46,8 @@
 public final class ActionEnvironment {
 
   /**
-   * A map of environment variables together with a list of variables to inherit from the shell
-   * environment.
+   * 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 {
 
@@ -78,14 +78,14 @@ default int 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() && base.isEmpty()) {
-        return EMPTY_ENVIRONMENT_VARIABLES;
+      if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
+        return base;
       }
       return new CompoundEnvironmentVariables(fixedVars, inheritedVars, base);
     }
@@ -168,8 +168,6 @@ public ImmutableSet<String> getInheritedEnvironment() {
    * 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()) {
@@ -190,9 +188,11 @@ private ActionEnvironment(EnvironmentVariables 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 vars) {
     if (vars.isEmpty()) {
@@ -201,6 +201,12 @@ private static ActionEnvironment create(EnvironmentVariables vars) {
     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 ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
@@ -214,21 +220,29 @@ public static ActionEnvironment create(Map<String, String> fixedEnv) {
    * 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> fixedVars) {
-    return ActionEnvironment.create(
-        CompoundEnvironmentVariables.create(fixedVars, ImmutableSet.of(), vars));
+  public ActionEnvironment withAdditionalFixedVariables(Map<String, String> fixedVars) {
+    return withAdditionalVariables(fixedVars, ImmutableSet.of());
   }
 
   /**
    * 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 addVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
-    return ActionEnvironment.create(
-        CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars));
+  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 the combined size of the fixed and inherited environments. */
+  /**
+   * 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 vars.size();
   }

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -134,7 +134,6 @@ java_library(
         ":test/execution_info",
         ":test/instrumented_files_info",
         ":test/test_configuration",
-        ":test/test_environment_info",
         ":test/test_sharding_strategy",
         ":test/test_trimming_transition_factory",
         ":toolchain_collection",
@@ -339,6 +338,7 @@ java_library(
         ":resolved_toolchain_context",
         ":rule_configured_object_value",
         ":rule_definition_environment",
+        ":run_environment_info",
         ":starlark/args",
         ":starlark/bazel_build_api_globals",
         ":starlark/function_transition_util",
@@ -356,7 +356,6 @@ java_library(
         ":test/execution_info",
         ":test/instrumented_files_info",
         ":test/test_configuration",
-        ":test/test_environment_info",
         ":test/test_sharding_strategy",
         ":toolchain_collection",
         ":toolchain_context",
@@ -999,6 +998,18 @@ java_library(
     ],
 )
 
+java_library(
+    name = "run_environment_info",
+    srcs = ["RunEnvironmentInfo.java"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/concurrent",
+        "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi",
+        "//src/main/java/net/starlark/java/eval",
+        "//third_party:guava",
+    ],
+)
+
 java_library(
     name = "rule_definition_environment",
     srcs = ["RuleDefinitionEnvironment.java"],
@@ -2365,17 +2376,6 @@ java_library(
     ],
 )
 
-java_library(
-    name = "test/test_environment_info",
-    srcs = ["test/TestEnvironmentInfo.java"],
-    deps = [
-        "//src/main/java/com/google/devtools/build/lib/concurrent",
-        "//src/main/java/com/google/devtools/build/lib/packages",
-        "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/test",
-        "//third_party:guava",
-    ],
-)
-
 java_library(
     name = "test/test_sharding_strategy",
     srcs = ["test/TestShardingStrategy.java"],

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

@@ -40,7 +40,6 @@
 import com.google.devtools.build.lib.analysis.test.InstrumentedFilesInfo;
 import com.google.devtools.build.lib.analysis.test.TestActionBuilder;
 import com.google.devtools.build.lib.analysis.test.TestConfiguration;
-import com.google.devtools.build.lib.analysis.test.TestEnvironmentInfo;
 import com.google.devtools.build.lib.analysis.test.TestProvider;
 import com.google.devtools.build.lib.analysis.test.TestProvider.TestParams;
 import com.google.devtools.build.lib.analysis.test.TestTagsProvider;
@@ -465,9 +464,8 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
                     providersBuilder.getProvider(
                         InstrumentedFilesInfo.STARLARK_CONSTRUCTOR.getKey()));
 
-    TestEnvironmentInfo environmentProvider =
-        (TestEnvironmentInfo)
-            providersBuilder.getProvider(TestEnvironmentInfo.PROVIDER.getKey());
+    RunEnvironmentInfo environmentProvider =
+        (RunEnvironmentInfo) providersBuilder.getProvider(RunEnvironmentInfo.PROVIDER.getKey());
     if (environmentProvider != null) {
       testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
       testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());

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

@@ -0,0 +1,106 @@
+// Copyright 2022 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.analysis;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.packages.BuiltinProvider;
+import com.google.devtools.build.lib.packages.NativeInfo;
+import com.google.devtools.build.lib.starlarkbuildapi.RunEnvironmentInfoApi;
+import java.util.List;
+import java.util.Map;
+import net.starlark.java.eval.Dict;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
+import net.starlark.java.eval.StarlarkList;
+
+/**
+ * Provider containing any additional environment variables for use in the executable or test
+ * action.
+ */
+@Immutable
+public final class RunEnvironmentInfo extends NativeInfo implements RunEnvironmentInfoApi {
+
+  /** Singleton instance of the provider type for {@link DefaultInfo}. */
+  public static final RunEnvironmentInfoProvider PROVIDER = new RunEnvironmentInfoProvider();
+
+  private final ImmutableMap<String, String> environment;
+  private final ImmutableList<String> inheritedEnvironment;
+  private final boolean shouldErrorOnNonExecutableRule;
+
+  /** Constructs a new provider with the given fixed and inherited environment variables. */
+  public RunEnvironmentInfo(
+      Map<String, String> environment,
+      List<String> inheritedEnvironment,
+      boolean shouldErrorOnNonExecutableRule) {
+    this.environment = ImmutableMap.copyOf(Preconditions.checkNotNull(environment));
+    this.inheritedEnvironment =
+        ImmutableList.copyOf(Preconditions.checkNotNull(inheritedEnvironment));
+    this.shouldErrorOnNonExecutableRule = shouldErrorOnNonExecutableRule;
+  }
+
+  @Override
+  public RunEnvironmentInfoProvider getProvider() {
+    return PROVIDER;
+  }
+
+  /**
+   * Returns environment variables which should be set when the target advertising this provider is
+   * executed.
+   */
+  @Override
+  public ImmutableMap<String, String> getEnvironment() {
+    return environment;
+  }
+
+  /**
+   * Returns names of environment variables whose value should be inherited from the shell
+   * environment when the target advertising this provider is executed.
+   */
+  @Override
+  public ImmutableList<String> getInheritedEnvironment() {
+    return inheritedEnvironment;
+  }
+
+  /**
+   * Returns whether advertising this provider on a non-executable (and thus non-test) rule should
+   * result in an error or a warning. The latter is required to not break testing.TestEnvironment,
+   * which is now implemented via RunEnvironmentInfo.
+   */
+  public boolean shouldErrorOnNonExecutableRule() {
+    return shouldErrorOnNonExecutableRule;
+  }
+
+  /** Provider implementation for {@link RunEnvironmentInfoApi}. */
+  public static class RunEnvironmentInfoProvider extends BuiltinProvider<RunEnvironmentInfo>
+      implements RunEnvironmentInfoApi.RunEnvironmentInfoApiProvider {
+
+    private RunEnvironmentInfoProvider() {
+      super("RunEnvironmentInfo", RunEnvironmentInfo.class);
+    }
+
+    @Override
+    public RunEnvironmentInfoApi constructor(
+        Dict<?, ?> environment, Sequence<?> inheritedEnvironment) throws EvalException {
+      return new RunEnvironmentInfo(
+          Dict.cast(environment, String.class, String.class, "environment"),
+          StarlarkList.immutableCopyOf(
+              Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")),
+          /* shouldErrorOnNonExecutableRule */ true);
+    }
+  }
+}

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

@@ -475,9 +475,7 @@ private static CommandLine computeArgs(RuleContext ruleContext, CommandLine addi
   }
 
   private static ActionEnvironment computeActionEnvironment(RuleContext ruleContext) {
-    // Currently, "env" and "env_inherit" are not added to Starlark-defined rules (unlike "args"),
-    // in order to avoid breaking existing Starlark rules that use those attribute names.
-    // TODO(brandjon): Support "env" and "env_inherit" for Starlark-defined rules.
+    // Executable Starlark rules can use RunEnvironmentInfo to specify environment variables.
     boolean isNativeRule =
         ruleContext.getRule().getRuleClassObject().getRuleDefinitionEnvironmentLabel() == null;
     if (!isNativeRule

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkModules.java

@@ -18,6 +18,7 @@
 import com.google.devtools.build.lib.analysis.ActionsProvider;
 import com.google.devtools.build.lib.analysis.DefaultInfo;
 import com.google.devtools.build.lib.analysis.OutputGroupInfo;
+import com.google.devtools.build.lib.analysis.RunEnvironmentInfo;
 import com.google.devtools.build.lib.packages.StarlarkLibrary;
 import com.google.devtools.build.lib.packages.StructProvider;
 import net.starlark.java.eval.Starlark;
@@ -44,5 +45,6 @@ public static void addPredeclared(ImmutableMap.Builder<String, Object> predeclar
     predeclared.put("OutputGroupInfo", OutputGroupInfo.STARLARK_CONSTRUCTOR);
     predeclared.put("Actions", ActionsProvider.INSTANCE);
     predeclared.put("DefaultInfo", DefaultInfo.PROVIDER);
+    predeclared.put("RunEnvironmentInfo", RunEnvironmentInfo.PROVIDER);
   }
 }

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleConfiguredTargetUtil.java

@@ -24,6 +24,7 @@
 import com.google.devtools.build.lib.analysis.DefaultInfo;
 import com.google.devtools.build.lib.analysis.RuleConfiguredTargetBuilder;
 import com.google.devtools.build.lib.analysis.RuleContext;
+import com.google.devtools.build.lib.analysis.RunEnvironmentInfo;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
@@ -348,6 +349,17 @@ private static void addProviders(
       if (getProviderKey(declaredProvider).equals(DefaultInfo.PROVIDER.getKey())) {
         parseDefaultProviderFields((DefaultInfo) declaredProvider, context, builder);
         defaultProviderProvidedExplicitly = true;
+      } else if (getProviderKey(declaredProvider).equals(RunEnvironmentInfo.PROVIDER.getKey())
+          && !(context.isExecutable() || context.getRuleContext().isTestTarget())) {
+        String message =
+            "Returning RunEnvironmentInfo from a non-executable, non-test target has no effect";
+        RunEnvironmentInfo runEnvironmentInfo = (RunEnvironmentInfo) declaredProvider;
+        if (runEnvironmentInfo.shouldErrorOnNonExecutableRule()) {
+          context.getRuleContext().ruleError(message);
+        } else {
+          context.getRuleContext().ruleWarning(message);
+          builder.addStarlarkDeclaredProvider(declaredProvider);
+        }
       } else {
         builder.addStarlarkDeclaredProvider(declaredProvider);
       }

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

@@ -453,7 +453,7 @@ private TestParams createTestAction(int shards)
                 testProperties,
                 runfilesSupport
                     .getActionEnvironment()
-                    .addVariables(extraTestEnv, extraInheritedEnv),
+                    .withAdditionalVariables(extraTestEnv, extraInheritedEnv),
                 executionSettings,
                 shard,
                 run,