Merge the action input prefetcher into the ExecutorBuilder.

This should be a no-op change, primarily intended to improve the BlazeModule
API. The code simplification in the distributor code path is incidental.

--
PiperOrigin-RevId: 144441458
MOS_MIGRATED_REVID=144441458

5 files changed, 36 insertions, 43 deletions

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 397521512e..c332127b8a 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,6 +64,7 @@ import com.google.devtools.build.lib.events.EventKind;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.exec.ActionContextConsumer;
 import com.google.devtools.build.lib.exec.ActionContextProvider;
+import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.BlazeExecutor;
 import com.google.devtools.build.lib.exec.CheckUpToDateFilter;
 import com.google.devtools.build.lib.exec.ExecutionOptions;
@@ -170,6 +171,7 @@ public class ExecutionTool {
   private final BuildRequest request;
   private BlazeExecutor executor;
   private final ActionInputFileCache fileCache;
+  private final ActionInputPrefetcher prefetcher;
   private final ImmutableList<ActionContextProvider> actionContextProviders;
 
   private Map<String, SpawnActionContext> spawnStrategyMap =
@@ -223,8 +225,13 @@ public class ExecutionTool {
               env.getExecRoot().getPathString(), env.getDirectories().getFileSystem());
     }
     this.fileCache = cache;
+    this.prefetcher = builder.getActionInputPrefetcher();
         
     this.actionContextProviders = builder.getActionContextProviders();
+    for (ActionContextProvider provider : actionContextProviders) {
+      provider.init(fileCache, prefetcher);
+    }
+
     StrategyConverter strategyConverter = new StrategyConverter(actionContextProviders);
 
     for (ActionContextConsumer consumer : builder.getActionContextConsumers()) {
@@ -398,10 +405,7 @@ public class ExecutionTool {
     boolean buildCompleted = false;
     try {
       for (ActionContextProvider actionContextProvider : actionContextProviders) {
-        actionContextProvider.executionPhaseStarting(
-            fileCache,
-            actionGraph,
-            allArtifactsForProviders);
+        actionContextProvider.executionPhaseStarting(actionGraph, allArtifactsForProviders);
       }
       executor.executionPhaseStarting();
       skyframeExecutor.drainChangedFiles();
diff --git a/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
index 6bdd81d56a..2362fcc847 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/ActionContextProvider.java
@@ -34,16 +34,25 @@ public abstract class ActionContextProvider {
   public abstract Iterable<ActionContext> getActionContexts();
 
   /**
+   * Two-phase initialization. The input file cache and the input prefetcher usually come from a
+   * different module than the {@link ActionContextProvider} instances that require them, so this
+   * method is called after {@link BlazeModule#executorInit}.
+   *
+   * @param actionInputFileCache the input file cache
+   * @param actionInputPrefetcher the input file prefetcher
+   */
+  public void init(
+      ActionInputFileCache actionInputFileCache, ActionInputPrefetcher actionInputPrefetcher) {
+  }
+
+  /**
    * Called when the executor is constructed. The parameter contains all the contexts that were
    * selected for this execution phase.
    */
   public void executorCreated(Iterable<ActionContext> usedContexts) throws ExecutorInitException {}
 
-  /**
-   * Called when the execution phase is started.
-   */
+  /** Called when the execution phase is started. */
   public void executionPhaseStarting(
-      ActionInputFileCache actionInputFileCache,
       ActionGraph actionGraph,
       Iterable<Artifact> topLevelArtifacts)
       throws ExecutorInitException, InterruptedException {}
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 5c90c6de49..6f9596d0a4 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
@@ -30,6 +30,7 @@ public class ExecutorBuilder {
   private final List<ActionContextProvider> actionContextProviders = new ArrayList<>();
   private final List<ActionContextConsumer> actionContextConsumers = new ArrayList<>();
   private ActionInputFileCache cache;
+  private ActionInputPrefetcher prefetcher;
 
   // These methods shouldn't be public, but they have to be right now as ExecutionTool is in another
   // package.
@@ -46,6 +47,10 @@ public class ExecutorBuilder {
     return cache;
   }
 
+  public ActionInputPrefetcher getActionInputPrefetcher() {
+    return prefetcher == null ? ActionInputPrefetcher.NONE : prefetcher;
+  }
+
   /**
    * Adds the specified action context providers to the executor.
    */
@@ -79,5 +84,14 @@ public class ExecutorBuilder {
     this.cache = Preconditions.checkNotNull(cache);
     return this;
   }
-}
 
+  /**
+   * Sets the action input prefetcher. Only one module may set the prefetcher. If multiple modules
+   * set it, this method will throw an {@link IllegalStateException}.
+   */
+  public ExecutorBuilder setActionInputPrefetcher(ActionInputPrefetcher prefetcher) {
+    Preconditions.checkState(this.prefetcher == null);
+    this.prefetcher = Preconditions.checkNotNull(prefetcher);
+    return this;
+  }
+}
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 dde71d48b8..39d6d4a9af 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
@@ -20,7 +20,6 @@ import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.ServerDirectories;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
 import com.google.devtools.build.lib.exec.OutputService;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
@@ -157,16 +156,6 @@ public abstract class BlazeModule {
   }
 
   /**
-   * Returns an implementation of {@link ActionInputPrefetcher}.
-   *
-   * <p>This method will be called at the beginning of each command (after #beforeCommand).
-   */
-  @SuppressWarnings("unused")
-  public ActionInputPrefetcher getPrefetcher() throws AbruptExitException {
-    return null;
-  }
-
-  /**
    * Does any handling of options needed by the command.
    *
    * <p>This method will be called at the beginning of each command (after #beforeCommand).
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java b/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
index cc473da5e6..29bf5a53c7 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
@@ -20,7 +20,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.eventbus.EventBus;
-import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.PackageRootResolver;
 import com.google.devtools.build.lib.actions.cache.ActionCache;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
@@ -33,7 +32,6 @@ import com.google.devtools.build.lib.analysis.config.DefaultsPackage;
 import com.google.devtools.build.lib.analysis.config.InvalidConfigurationException;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Reporter;
-import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.OutputService;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
 import com.google.devtools.build.lib.packages.Target;
@@ -92,7 +90,6 @@ public final class CommandEnvironment {
   private PathFragment relativeWorkingDirectory = PathFragment.EMPTY_FRAGMENT;
   private long commandStartTime;
   private OutputService outputService;
-  private ActionInputPrefetcher actionInputPrefetcher;
   private Path workingDirectory;
 
   private String commandName;
@@ -346,10 +343,6 @@ public final class CommandEnvironment {
     return outputService;
   }
 
-  public ActionInputPrefetcher getActionInputPrefetcher() {
-    return actionInputPrefetcher == null ? ActionInputPrefetcher.NONE : actionInputPrefetcher;
-  }
-
   public ActionCache getPersistentActionCache() throws IOException {
     return workspace.getPersistentActionCache(reporter);
   }
@@ -535,7 +528,6 @@ public final class CommandEnvironment {
 
     outputService = null;
     BlazeModule outputModule = null;
-    ImmutableList.Builder<ActionInputPrefetcher> prefetchersBuilder = ImmutableList.builder();
     if (command.builds()) {
       for (BlazeModule module : runtime.getBlazeModules()) {
         OutputService moduleService = module.getOutputService();
@@ -549,23 +541,8 @@ public final class CommandEnvironment {
           outputService = moduleService;
           outputModule = module;
         }
-
-        ActionInputPrefetcher actionInputPrefetcher = module.getPrefetcher();
-        if (actionInputPrefetcher != null) {
-          prefetchersBuilder.add(actionInputPrefetcher);
-        }
       }
     }
-    final ImmutableList<ActionInputPrefetcher> actionInputPrefetchers = prefetchersBuilder.build();
-    actionInputPrefetcher =
-        new ActionInputPrefetcher() {
-          @Override
-          public void prefetchFile(ActionInput input) {
-            for (ActionInputPrefetcher prefetcher : actionInputPrefetchers) {
-              prefetcher.prefetchFile(input);
-            }
-          }
-        };
 
     SkyframeExecutor skyframeExecutor = getSkyframeExecutor();
     skyframeExecutor.setOutputService(outputService);