Update DynamicExecutionModule to use registry methods instead of executorInit where possible.
Part of the roll-forward of https://github.com/bazelbuild/bazel/commit/37aeabcd39fe326d1c4e55693d8d207f9f7ac6c4.
PiperOrigin-RevId: 303951192
diff --git a/src/main/java/com/google/devtools/build/lib/dynamic/DynamicExecutionModule.java b/src/main/java/com/google/devtools/build/lib/dynamic/DynamicExecutionModule.java
index d7126c8..6f3bc93 100644
--- a/src/main/java/com/google/devtools/build/lib/dynamic/DynamicExecutionModule.java
+++ b/src/main/java/com/google/devtools/build/lib/dynamic/DynamicExecutionModule.java
@@ -24,10 +24,9 @@
import com.google.devtools.build.lib.actions.Spawn;
import com.google.devtools.build.lib.actions.SpawnStrategy;
import com.google.devtools.build.lib.actions.Spawns;
-import com.google.devtools.build.lib.buildtool.BuildRequest;
import com.google.devtools.build.lib.concurrent.ExecutorUtil;
import com.google.devtools.build.lib.exec.ExecutionPolicy;
-import com.google.devtools.build.lib.exec.ExecutorBuilder;
+import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
import com.google.devtools.build.lib.runtime.BlazeModule;
import com.google.devtools.build.lib.runtime.Command;
import com.google.devtools.build.lib.runtime.CommandEnvironment;
@@ -42,9 +41,6 @@
*/
public class DynamicExecutionModule extends BlazeModule {
- /** Strings that can be used to select this strategy in flags. */
- private static final String[] COMMANDLINE_IDENTIFIERS = {"dynamic", "dynamic_worker"};
-
private ExecutorService executorService;
public DynamicExecutionModule() {}
@@ -105,47 +101,46 @@
: options.dynamicRemoteStrategy;
}
+ @Override
+ public void registerSpawnStrategies(
+ SpawnStrategyRegistry.Builder registryBuilder, CommandEnvironment env)
+ throws ExecutorInitException {
+ registerSpawnStrategies(
+ registryBuilder, env.getOptions().getOptions(DynamicExecutionOptions.class));
+ }
+
+ // CommandEnvironment is difficult to access in tests, so use this method for testing.
@VisibleForTesting
- void initStrategies(ExecutorBuilder builder, DynamicExecutionOptions options)
+ final void registerSpawnStrategies(
+ SpawnStrategyRegistry.Builder registryBuilder, DynamicExecutionOptions options)
throws ExecutorInitException {
if (!options.internalSpawnScheduler) {
return;
}
+ SpawnStrategy strategy;
if (options.legacySpawnScheduler) {
- builder.addActionContext(
- SpawnStrategy.class,
- new LegacyDynamicSpawnStrategy(executorService, options, this::getExecutionPolicy),
- COMMANDLINE_IDENTIFIERS);
+ strategy = new LegacyDynamicSpawnStrategy(executorService, options, this::getExecutionPolicy);
} else {
- builder.addActionContext(
- SpawnStrategy.class,
- new DynamicSpawnStrategy(executorService, options, this::getExecutionPolicy),
- COMMANDLINE_IDENTIFIERS);
+ strategy = new DynamicSpawnStrategy(executorService, options, this::getExecutionPolicy);
}
- builder.addStrategyByContext(SpawnStrategy.class, "dynamic");
+ registryBuilder.registerStrategy(strategy, "dynamic", "dynamic_worker");
for (Map.Entry<String, List<String>> mnemonicToStrategies : getLocalStrategies(options)) {
throwIfContainsDynamic(mnemonicToStrategies.getValue(), "--dynamic_local_strategy");
- builder.addDynamicLocalStrategiesByMnemonic(
+ registryBuilder.addDynamicLocalStrategiesByMnemonic(
mnemonicToStrategies.getKey(), mnemonicToStrategies.getValue());
}
for (Map.Entry<String, List<String>> mnemonicToStrategies : getRemoteStrategies(options)) {
throwIfContainsDynamic(mnemonicToStrategies.getValue(), "--dynamic_remote_strategy");
- builder.addDynamicRemoteStrategiesByMnemonic(
+ registryBuilder.addDynamicRemoteStrategiesByMnemonic(
mnemonicToStrategies.getKey(), mnemonicToStrategies.getValue());
}
}
- @Override
- public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorBuilder builder)
- throws ExecutorInitException {
- initStrategies(builder, env.getOptions().getOptions(DynamicExecutionOptions.class));
- }
-
private void throwIfContainsDynamic(List<String> strategies, String flagName)
throws ExecutorInitException {
- ImmutableSet<String> identifiers = ImmutableSet.copyOf(COMMANDLINE_IDENTIFIERS);
+ ImmutableSet<String> identifiers = ImmutableSet.of("dynamic", "dynamic_worker");
if (!Sets.intersection(identifiers, ImmutableSet.copyOf(strategies)).isEmpty()) {
throw new ExecutorInitException(
"Cannot use strategy "
@@ -158,6 +153,7 @@
/**
* Use the {@link Spawn} metadata to determine if it can be executed locally, remotely, or both.
+ *
* @param spawn the {@link Spawn} action
* @return the {@link ExecutionPolicy} containing local/remote execution policies
*/
diff --git a/src/test/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategyTest.java b/src/test/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategyTest.java
index 78aa053..39a2f6f 100644
--- a/src/test/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategyTest.java
@@ -48,6 +48,7 @@
import com.google.devtools.build.lib.exec.ExecutionOptions;
import com.google.devtools.build.lib.exec.ExecutorBuilder;
import com.google.devtools.build.lib.exec.SpawnActionContextMaps;
+import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
import com.google.devtools.build.lib.testutil.TestThread;
import com.google.devtools.build.lib.testutil.TestUtils;
import com.google.devtools.build.lib.util.io.FileOutErr;
@@ -316,16 +317,22 @@
checkState(executorServiceForCleanup == null);
executorServiceForCleanup = executorService;
- ExecutorBuilder executorBuilder =
- new ExecutorBuilder()
- .addActionContext(SpawnStrategy.class, localStrategy, "mock-local")
- .addActionContext(SpawnStrategy.class, remoteStrategy, "mock-remote");
+ ExecutorBuilder executorBuilder = new ExecutorBuilder();
+ SpawnStrategyRegistry.Builder spawnStrategyRegistryBuilder =
+ executorBuilder.asSpawnStrategyRegistryBuilder();
+
+ spawnStrategyRegistryBuilder.registerStrategy(localStrategy, "mock-local");
+ spawnStrategyRegistryBuilder.registerStrategy(remoteStrategy, "mock-remote");
if (sandboxedStrategy != null) {
- executorBuilder.addActionContext(SpawnStrategy.class, sandboxedStrategy, "mock-sandboxed");
+ spawnStrategyRegistryBuilder.registerStrategy(sandboxedStrategy, "mock-sandboxed");
}
- new DynamicExecutionModule(executorService).initStrategies(executorBuilder, options);
+ DynamicExecutionModule dynamicExecutionModule = new DynamicExecutionModule(executorService);
+ // TODO(b/63987502): This only exists during the migration to SpawnStrategyRegistry.
+ executorBuilder.addStrategyByContext(SpawnStrategy.class, "dynamic");
+ dynamicExecutionModule.registerSpawnStrategies(spawnStrategyRegistryBuilder, options);
+
SpawnActionContextMaps spawnActionContextMaps = executorBuilder.getSpawnActionContextMaps();
Executor executor =
