Let Starlark executable rules specify their environment

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

@@ -133,7 +133,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",
@@ -340,6 +339,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",
@@ -357,7 +357,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",
@@ -1003,6 +1002,19 @@ 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/com/google/devtools/build/lib/starlarkbuildapi/test",
+        "//src/main/java/net/starlark/java/eval",
+        "//third_party:guava",
+    ],
+)
+
 java_library(
     name = "rule_definition_environment",
     srcs = ["RuleDefinitionEnvironment.java"],
@@ -2411,17 +2423,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;
@@ -470,9 +469,9 @@ 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,92 @@
+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;
+
+@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 Map<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(b/176554800): 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,15 @@ 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/rules/BUILD

@@ -123,12 +123,12 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
+        "//src/main/java/com/google/devtools/build/lib/analysis:run_environment_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/analysis_failure_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/analysis_test_result_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/execution_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/instrumented_files_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/test_configuration",
-        "//src/main/java/com/google/devtools/build/lib/analysis:test/test_environment_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_provider",
         "//src/main/java/com/google/devtools/build/lib/concurrent",

src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java

@@ -13,8 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.rules.test;
 
+import com.google.devtools.build.lib.analysis.RunEnvironmentInfo;
 import com.google.devtools.build.lib.analysis.test.ExecutionInfo;
-import com.google.devtools.build.lib.analysis.test.TestEnvironmentInfo;
 import com.google.devtools.build.lib.starlarkbuildapi.test.TestingModuleApi;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
@@ -31,13 +31,14 @@ public ExecutionInfo executionInfo(Dict<?, ?> requirements /* <String, String> *
   }
 
   @Override
-  public TestEnvironmentInfo testEnvironment(
+  public RunEnvironmentInfo testEnvironment(
       Dict<?, ?> environment /* <String, String> */,
       Sequence<?> inheritedEnvironment /* <String> */)
       throws EvalException {
-    return new TestEnvironmentInfo(
+    return new RunEnvironmentInfo(
         Dict.cast(environment, String.class, String.class, "environment"),
         StarlarkList.immutableCopyOf(
-            Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")));
+            Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")),
+        /* shouldErrorOnNonExecutableRule */ false);
   }
 }

src/main/java/com/google/devtools/build/lib/runtime/commands/BUILD

@@ -47,6 +47,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:no_build_event",
         "//src/main/java/com/google/devtools/build/lib/analysis:no_build_request_finished_event",
         "//src/main/java/com/google/devtools/build/lib/analysis:print_action_visitor",
+        "//src/main/java/com/google/devtools/build/lib/analysis:run_environment_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:test/test_configuration",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream",
         "//src/main/java/com/google/devtools/build/lib/cmdline",

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

@@ -20,9 +20,11 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.devtools.build.lib.actions.ActionEnvironment;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.CommandLine;
@@ -32,6 +34,7 @@
 import com.google.devtools.build.lib.analysis.AliasProvider;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
+import com.google.devtools.build.lib.analysis.RunEnvironmentInfo;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.analysis.ShToolchain;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
@@ -487,11 +490,18 @@ public BlazeCommandResult exec(CommandEnvironment env, OptionsParsingResult opti
     } else {
       workingDir =
           targetToRunRunfilesDir != null ? targetToRunRunfilesDir : env.getWorkingDirectory();
+      ActionEnvironment actionEnvironment = ActionEnvironment.EMPTY;
       if (targetToRunRunfilesSupport != null) {
-        targetToRunRunfilesSupport
-            .getActionEnvironment()
-            .resolve(runEnvironment, env.getClientEnv());
+        actionEnvironment = targetToRunRunfilesSupport.getActionEnvironment();
       }
+      RunEnvironmentInfo environmentProvider = targetToRun.get(RunEnvironmentInfo.PROVIDER);
+      if (environmentProvider != null) {
+        actionEnvironment = actionEnvironment.withAdditionalVariables(
+            environmentProvider.getEnvironment(),
+            ImmutableSet.copyOf(environmentProvider.getInheritedEnvironment())
+        );
+      }
+      actionEnvironment.resolve(runEnvironment, env.getClientEnv());
       try {
         List<String> args = computeArgs(targetToRun, commandLineArgs);
         constructCommandLine(