Add registerActionContexts and registerSpawnStrategies methods to BlazeModule.

Part of the rollforward of https://github.com/bazelbuild/bazel/commit/37aeabcd39fe326d1c4e55693d8d207f9f7ac6c4.

PiperOrigin-RevId: 303378498
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
index 3eea94d..5cffb2e 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -64,7 +64,9 @@
 import com.google.devtools.build.lib.exec.ExecutionOptions;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
 import com.google.devtools.build.lib.exec.ExecutorLifecycleListener;
+import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
 import com.google.devtools.build.lib.exec.SpawnActionContextMaps;
+import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
 import com.google.devtools.build.lib.exec.SymlinkTreeStrategy;
 import com.google.devtools.build.lib.packages.StarlarkSemanticsOptions;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
@@ -136,24 +138,39 @@
       throw new ExecutorInitException("Execroot creation failed", e);
     }
 
-    ExecutorBuilder builder = new ExecutorBuilder();
+    ExecutorBuilder executorBuilder = new ExecutorBuilder();
+    ModuleActionContextRegistry.Builder actionContextRegistryBuilder =
+        executorBuilder.asModuleActionContextRegistryBuilder();
+    SpawnStrategyRegistry.Builder spawnStrategyRegistryBuilder =
+        executorBuilder.asSpawnStrategyRegistryBuilder();
+
     for (BlazeModule module : runtime.getBlazeModules()) {
-      try (SilentCloseable closeable = Profiler.instance().profile(module + ".executorInit")) {
-        module.executorInit(env, request, builder);
+      try (SilentCloseable ignored = Profiler.instance().profile(module + ".executorInit")) {
+        module.executorInit(env, request, executorBuilder);
+      }
+
+      try (SilentCloseable ignored =
+          Profiler.instance().profile(module + ".registerActionContexts")) {
+        module.registerActionContexts(actionContextRegistryBuilder, env, request);
+      }
+
+      try (SilentCloseable ignored =
+          Profiler.instance().profile(module + ".registerSpawnStrategies")) {
+        module.registerSpawnStrategies(spawnStrategyRegistryBuilder, env);
       }
     }
-    builder.addActionContext(
+    actionContextRegistryBuilder.register(
         SymlinkTreeActionContext.class,
         new SymlinkTreeStrategy(env.getOutputService(), env.getBlazeWorkspace().getBinTools()));
     // TODO(philwo) - the ExecutionTool should not add arbitrary dependencies on its own, instead
     // these dependencies should be added to the ActionContextConsumer of the module that actually
     // depends on them.
-    builder
-        .addStrategyByContext(WorkspaceStatusAction.Context.class, "")
-        .addStrategyByContext(SymlinkTreeActionContext.class, "");
+    actionContextRegistryBuilder
+        .restrictTo(WorkspaceStatusAction.Context.class, "")
+        .restrictTo(SymlinkTreeActionContext.class, "");
 
-    this.prefetcher = builder.getActionInputPrefetcher();
-    this.executorLifecycleListeners = builder.getExecutorLifecycleListeners();
+    this.prefetcher = executorBuilder.getActionInputPrefetcher();
+    this.executorLifecycleListeners = executorBuilder.getExecutorLifecycleListeners();
 
     // There are many different SpawnActions, and we want to control the action context they use
     // independently from each other, for example, to run genrules locally and Java compile action
@@ -161,8 +178,9 @@
     // context class, but also the mnemonic of the action.
     ExecutionOptions options = request.getOptions(ExecutionOptions.class);
     // TODO(jmmv): This should live in some testing-related Blaze module, not here.
-    builder.addStrategyByContext(TestActionContext.class, options.testStrategy);
-    spawnActionContextMaps = builder.getSpawnActionContextMaps();
+    actionContextRegistryBuilder.restrictTo(TestActionContext.class, options.testStrategy);
+
+    spawnActionContextMaps = executorBuilder.getSpawnActionContextMaps();
 
     if (options.availableResources != null && options.removeLocalResources) {
       throw new ExecutorInitException(
diff --git a/src/main/java/com/google/devtools/build/lib/exec/ExecutorBuilder.java b/src/main/java/com/google/devtools/build/lib/exec/ExecutorBuilder.java
index 73da12989..ab21e51 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/ExecutorBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/ExecutorBuilder.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.exec;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
@@ -161,4 +162,120 @@
     executorLifecycleListeners.add(listener);
     return this;
   }
+
+  // TODO(katre): Use a fake implementation to allow for migration to the new API.
+  public ModuleActionContextRegistry.Builder asModuleActionContextRegistryBuilder() {
+    return new ModuleActionContextDelegate(this);
+  }
+
+  private static final class ModuleActionContextDelegate
+      implements ModuleActionContextRegistry.Builder {
+    private final ExecutorBuilder executorBuilder;
+
+    private ModuleActionContextDelegate(ExecutorBuilder executorBuilder) {
+      this.executorBuilder = executorBuilder;
+    }
+
+    @Override
+    public ModuleActionContextRegistry.Builder restrictTo(
+        Class<?> identifyingType, String restriction) {
+      Preconditions.checkArgument(ActionContext.class.isAssignableFrom(identifyingType));
+      @SuppressWarnings("unchecked")
+      Class<? extends ActionContext> castType = (Class<? extends ActionContext>) identifyingType;
+      this.executorBuilder.addStrategyByContext(castType, restriction);
+      return this;
+    }
+
+    @Override
+    public <T extends ActionContext> ModuleActionContextRegistry.Builder register(
+        Class<T> identifyingType, T context, String... commandLineIdentifiers) {
+      this.executorBuilder.addActionContext(identifyingType, context, commandLineIdentifiers);
+      return this;
+    }
+
+    @Override
+    public ModuleActionContextRegistry build() throws ExecutorInitException {
+      throw new UnsupportedOperationException("not a real builder");
+    }
+  }
+
+  // TODO(katre): Use a fake implementation to allow for migration to the new API.
+  public SpawnStrategyRegistry.Builder asSpawnStrategyRegistryBuilder() {
+    return new SpawnStrategyRegistryDelegate(this);
+  }
+
+  private static final class SpawnStrategyRegistryDelegate
+      implements SpawnStrategyRegistry.Builder {
+    private final ExecutorBuilder executorBuilder;
+
+    private SpawnStrategyRegistryDelegate(ExecutorBuilder executorBuilder) {
+      this.executorBuilder = executorBuilder;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder addDescriptionFilter(
+        RegexFilter filter, List<String> identifiers) {
+      this.executorBuilder.addStrategyByRegexp(filter, identifiers);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder addMnemonicFilter(
+        String mnemonic, List<String> identifiers) {
+      this.executorBuilder.addStrategyByMnemonic(mnemonic, identifiers);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder registerStrategy(
+        SpawnStrategy strategy, List<String> commandlineIdentifiers) {
+      this.executorBuilder.addActionContext(
+          SpawnStrategy.class, strategy, commandlineIdentifiers.toArray(new String[0]));
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder useLegacyDescriptionFilterPrecedence() {
+      // Ignored.
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder setDefaultStrategies(List<String> defaultStrategies) {
+      this.executorBuilder.addStrategyByMnemonic("", defaultStrategies);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder resetDefaultStrategies() {
+      this.executorBuilder.addStrategyByMnemonic("", ImmutableList.of(""));
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder addDynamicRemoteStrategiesByMnemonic(
+        String mnemonic, List<String> strategies) {
+      this.executorBuilder.addDynamicRemoteStrategiesByMnemonic(mnemonic, strategies);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder addDynamicLocalStrategiesByMnemonic(
+        String mnemonic, List<String> strategies) {
+      this.executorBuilder.addDynamicLocalStrategiesByMnemonic(mnemonic, strategies);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry.Builder setRemoteLocalFallbackStrategyIdentifier(
+        String commandlineIdentifier) {
+      this.executorBuilder.setRemoteFallbackStrategy(commandlineIdentifier);
+      return this;
+    }
+
+    @Override
+    public SpawnStrategyRegistry build() throws ExecutorInitException {
+      throw new UnsupportedOperationException("not a real builder");
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
index 54f64ea..4938c3e 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.eventbus.SubscriberExceptionContext;
 import com.google.common.eventbus.SubscriberExceptionHandler;
+import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.BlazeVersionInfo;
@@ -31,6 +32,8 @@
 import com.google.devtools.build.lib.clock.Clock;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
+import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
+import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.PackageFactory;
 import com.google.devtools.build.lib.packages.PackageValidator;
@@ -306,6 +309,37 @@
       throws ExecutorInitException {}
 
   /**
+   * Registers any action contexts this module provides with the execution phase. They will be
+   * available for {@linkplain ActionContext.ActionContextRegistry#getContext querying} to actions
+   * and other action contexts.
+   *
+   * <p>This method is invoked before actions are executed but after {@link #executorInit}.
+   *
+   * @param registryBuilder builder with which to register action contexts
+   * @param env environment for the current command
+   * @param buildRequest the current build request
+   * @throws ExecutorInitException if there are fatal issues creating or registering action contexts
+   */
+  public void registerActionContexts(
+      ModuleActionContextRegistry.Builder registryBuilder,
+      CommandEnvironment env,
+      BuildRequest buildRequest)
+      throws ExecutorInitException {}
+
+  /**
+   * Registers any spawn strategies this module provides with the execution phase.
+   *
+   * <p>This method is invoked before actions are executed but after {@link #executorInit}.
+   *
+   * @param registryBuilder builder with which to register strategies
+   * @param env environment for the current command
+   * @throws ExecutorInitException if there are fatal issues creating or registering strategies
+   */
+  public void registerSpawnStrategies(
+      SpawnStrategyRegistry.Builder registryBuilder, CommandEnvironment env)
+      throws ExecutorInitException {}
+
+  /**
    * Called after each command.
    *
    * @throws AbruptExitException modules can throw this exception to modify the command exit code