Update RemoteModule to use registry methods instead of executorInit where possible.
Part of the roll-forward of https://github.com/bazelbuild/bazel/commit/37aeabcd39fe326d1c4e55693d8d207f9f7ac6c4.
PiperOrigin-RevId: 303938923
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
index 0a0ed81..3506993 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
@@ -22,12 +22,12 @@
import com.google.devtools.build.lib.actions.ActionGraph;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ExecutorInitException;
-import com.google.devtools.build.lib.actions.SpawnStrategy;
import com.google.devtools.build.lib.analysis.ArtifactsToOwnerLabels;
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.SpawnCache;
+import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
import com.google.devtools.build.lib.remote.options.RemoteOptions;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.runtime.CommandEnvironment;
@@ -80,45 +80,59 @@
env, cache, executor, retryScheduler, digestUtil, logDir);
}
- /** Registers the action contexts whose lifecycle this class manages. */
- public void registerActionContexts(ExecutorBuilder executorBuilder) {
- ExecutionOptions executionOptions =
- checkNotNull(env.getOptions().getOptions(ExecutionOptions.class));
- RemoteOptions remoteOptions = checkNotNull(env.getOptions().getOptions(RemoteOptions.class));
- String buildRequestId = env.getBuildRequestId();
- String commandId = env.getCommandId().toString();
-
+ /**
+ * Registers a remote spawn strategy if this instance was created with an executor, otherwise does
+ * nothing.
+ *
+ * @param registryBuilder builder with which to register the strategy
+ */
+ public void registerRemoteSpawnStrategyIfApplicable(
+ SpawnStrategyRegistry.Builder registryBuilder) {
if (executor == null) {
- RemoteSpawnCache spawnCache =
- new RemoteSpawnCache(
- env.getExecRoot(),
- remoteOptions,
- cache,
- buildRequestId,
- commandId,
- env.getReporter(),
- digestUtil,
- filesToDownload);
- executorBuilder.addActionContext(SpawnCache.class, spawnCache, "remote-cache");
- } else {
- RemoteSpawnRunner spawnRunner =
- new RemoteSpawnRunner(
- env.getExecRoot(),
- remoteOptions,
- env.getOptions().getOptions(ExecutionOptions.class),
- executionOptions.verboseFailures,
- env.getReporter(),
- buildRequestId,
- commandId,
- (RemoteExecutionCache) cache,
- executor,
- retryScheduler,
- digestUtil,
- logDir,
- filesToDownload);
- executorBuilder.addActionContext(
- SpawnStrategy.class, new RemoteSpawnStrategy(env.getExecRoot(), spawnRunner), "remote");
+ return; // Can't use a spawn strategy without executor.
}
+
+ RemoteSpawnRunner spawnRunner =
+ new RemoteSpawnRunner(
+ env.getExecRoot(),
+ checkNotNull(env.getOptions().getOptions(RemoteOptions.class)),
+ env.getOptions().getOptions(ExecutionOptions.class),
+ checkNotNull(env.getOptions().getOptions(ExecutionOptions.class)).verboseFailures,
+ env.getReporter(),
+ env.getBuildRequestId(),
+ env.getCommandId().toString(),
+ (RemoteExecutionCache) cache,
+ executor,
+ retryScheduler,
+ digestUtil,
+ logDir,
+ filesToDownload);
+ registryBuilder.registerStrategy(
+ new RemoteSpawnStrategy(env.getExecRoot(), spawnRunner), "remote");
+ }
+
+ /**
+ * Registers a spawn cache action context if this instance was created without an executor,
+ * otherwise does nothing.
+ *
+ * @param registryBuilder builder with which to register the cache
+ */
+ public void registerSpawnCacheIfApplicable(ModuleActionContextRegistry.Builder registryBuilder) {
+ if (executor != null) {
+ return; // No need to register cache if we're using a remote executor.
+ }
+
+ RemoteSpawnCache spawnCache =
+ new RemoteSpawnCache(
+ env.getExecRoot(),
+ checkNotNull(env.getOptions().getOptions(RemoteOptions.class)),
+ cache,
+ env.getBuildRequestId(),
+ env.getCommandId().toString(),
+ env.getReporter(),
+ digestUtil,
+ filesToDownload);
+ registryBuilder.register(SpawnCache.class, spawnCache, "remote-cache");
}
/** Returns the remote cache. */
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
index 1926d3d..b834667 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
@@ -43,6 +43,8 @@
import com.google.devtools.build.lib.events.Event;
import com.google.devtools.build.lib.events.Reporter;
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.TargetUtils;
import com.google.devtools.build.lib.remote.common.RemoteCacheClient;
import com.google.devtools.build.lib.remote.downloader.GrpcRemoteDownloader;
@@ -567,6 +569,31 @@
}
@Override
+ public void registerSpawnStrategies(
+ SpawnStrategyRegistry.Builder registryBuilder, CommandEnvironment env) {
+ if (actionContextProvider == null) {
+ return;
+ }
+ RemoteOptions remoteOptions =
+ Preconditions.checkNotNull(
+ env.getOptions().getOptions(RemoteOptions.class), "RemoteOptions");
+ registryBuilder.setRemoteLocalFallbackStrategyIdentifier(
+ remoteOptions.remoteLocalFallbackStrategy);
+ actionContextProvider.registerRemoteSpawnStrategyIfApplicable(registryBuilder);
+ }
+
+ @Override
+ public void registerActionContexts(
+ ModuleActionContextRegistry.Builder registryBuilder,
+ CommandEnvironment env,
+ BuildRequest buildRequest) {
+ if (actionContextProvider == null) {
+ return;
+ }
+ actionContextProvider.registerSpawnCacheIfApplicable(registryBuilder);
+ }
+
+ @Override
public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorBuilder builder) {
Preconditions.checkState(actionInputFetcher == null, "actionInputFetcher must be null");
Preconditions.checkNotNull(remoteOutputsMode, "remoteOutputsMode must not be null");
@@ -574,7 +601,6 @@
if (actionContextProvider == null) {
return;
}
- actionContextProvider.registerActionContexts(builder);
builder.addExecutorLifecycleListener(actionContextProvider);
RemoteOptions remoteOptions =
Preconditions.checkNotNull(
@@ -590,8 +616,6 @@
builder.setActionInputPrefetcher(actionInputFetcher);
remoteOutputService.setActionInputFetcher(actionInputFetcher);
}
-
- builder.setRemoteFallbackStrategy(remoteOptions.remoteLocalFallbackStrategy);
}
@Override
