Let Starlark executable rules specify their environment

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,