Unify input prefetching

All prefetching now goes through AbstractSpawnStrategy's implementation of
SpawnExecutionPolicy. Make sure the sandbox runners also do this consistently.

PiperOrigin-RevId: 164836877

6 files changed, 31 insertions, 22 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 953ae1647b..ed15ed5d17 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
@@ -229,7 +229,7 @@ public class ExecutionTool {
         
     this.actionContextProviders = builder.getActionContextProviders();
     for (ActionContextProvider provider : actionContextProviders) {
-      provider.init(fileCache, prefetcher);
+      provider.init(fileCache);
     }
 
     StrategyConverter strategyConverter = new StrategyConverter(actionContextProviders);
diff --git a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java
index ffef030501..5aacd8ab67 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java
@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.exec;
 
+import com.google.common.base.Predicates;
+import com.google.common.collect.Iterables;
 import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionInput;
@@ -72,6 +74,9 @@ public abstract class AbstractSpawnStrategy implements SandboxedSpawnActionConte
     final Duration timeout = Spawns.getTimeout(spawn);
     SpawnExecutionPolicy policy = new SpawnExecutionPolicy() {
       private final int id = execCount.incrementAndGet();
+      // Memoize the input mapping so that prefetchInputs can reuse it instead of recomputing it.
+      // TODO(ulfjack): Guard against client modification of this map.
+      private SortedMap<PathFragment, ActionInput> lazyInputMapping;
 
       @Override
       public int getId() {
@@ -79,8 +84,14 @@ public abstract class AbstractSpawnStrategy implements SandboxedSpawnActionConte
       }
 
       @Override
-      public void prefetchInputs(Iterable<ActionInput> inputs) throws IOException {
-        actionExecutionContext.getActionInputPrefetcher().prefetchFiles(inputs);
+      public void prefetchInputs() throws IOException {
+        if (Spawns.shouldPrefetchInputsForLocalExecution(spawn)) {
+          // TODO(philwo): Benchmark whether using an ExecutionService to do multiple operations in
+          // parallel speeds up prefetching of inputs.
+          // TODO(philwo): Do we have to expand middleman artifacts here?
+          actionExecutionContext.getActionInputPrefetcher().prefetchFiles(
+              Iterables.filter(getInputMapping().values(), Predicates.notNull()));
+        }
       }
 
       @Override
@@ -115,11 +126,14 @@ public abstract class AbstractSpawnStrategy implements SandboxedSpawnActionConte
 
       @Override
       public SortedMap<PathFragment, ActionInput> getInputMapping() throws IOException {
-        return spawnInputExpander.getInputMapping(
-            spawn,
-            actionExecutionContext.getArtifactExpander(),
-            actionExecutionContext.getActionInputFileCache(),
-            actionExecutionContext.getContext(FilesetActionContext.class));
+        if (lazyInputMapping == null) {
+          lazyInputMapping = spawnInputExpander.getInputMapping(
+              spawn,
+              actionExecutionContext.getArtifactExpander(),
+              actionExecutionContext.getActionInputFileCache(),
+              actionExecutionContext.getContext(FilesetActionContext.class));
+        }
+        return lazyInputMapping;
       }
 
       @Override
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 27e147b701..f2f5e0f6fb 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
@@ -16,7 +16,6 @@ package com.google.devtools.build.lib.exec;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionGraph;
 import com.google.devtools.build.lib.actions.ActionInputFileCache;
-import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
 
@@ -35,15 +34,13 @@ public abstract class ActionContextProvider {
   public abstract Iterable<? extends 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 com.google.devtools.build.lib.runtime.BlazeModule#executorInit}.
+   * Two-phase initialization. The input file cache usually comes from a different module than the
+   * {@link ActionContextProvider} instances that require it, so this method is called after
+   * {@link com.google.devtools.build.lib.runtime.BlazeModule#executorInit}.
    *
    * @param actionInputFileCache the input file cache
-   * @param actionInputPrefetcher the input file prefetcher
    */
-  public void init(
-      ActionInputFileCache actionInputFileCache, ActionInputPrefetcher actionInputPrefetcher) {
+  public void init(ActionInputFileCache actionInputFileCache) {
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java b/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java
index b98ca2ee59..e7f67042f0 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java
@@ -147,7 +147,7 @@ public interface SpawnRunner {
      * again. I suppose we could require implementations to memoize getInputMapping (but not compute
      * it eagerly), and that may change in the future.
      */
-    void prefetchInputs(Iterable<ActionInput> inputs) throws IOException;
+    void prefetchInputs() throws IOException;
 
     /**
      * The input file metadata cache for this specific spawn, which can be used to efficiently
diff --git a/src/main/java/com/google/devtools/build/lib/exec/local/LocalSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/exec/local/LocalSpawnRunner.java
index 548b3b4091..57d757691d 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/local/LocalSpawnRunner.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/local/LocalSpawnRunner.java
@@ -18,8 +18,6 @@ import static java.util.logging.Level.INFO;
 import static java.util.logging.Level.SEVERE;
 
 import com.google.common.base.Joiner;
-import com.google.common.base.Predicates;
-import com.google.common.collect.Iterables;
 import com.google.common.io.ByteStreams;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ResourceManager;
@@ -230,8 +228,7 @@ public final class LocalSpawnRunner implements SpawnRunner {
       if (Spawns.shouldPrefetchInputsForLocalExecution(spawn)) {
         stepLog(INFO, "prefetching inputs for local execution");
         setState(State.PREFETCHING_LOCAL_INPUTS);
-        policy.prefetchInputs(
-            Iterables.filter(policy.getInputMapping().values(), Predicates.notNull()));
+        policy.prefetchInputs();
       }
 
       stepLog(INFO, "running locally");
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/AbstractSandboxSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/sandbox/AbstractSandboxSpawnRunner.java
index 7fb7274e30..6ac4a6e319 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/AbstractSandboxSpawnRunner.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/AbstractSandboxSpawnRunner.java
@@ -90,10 +90,11 @@ abstract class AbstractSandboxSpawnRunner implements SpawnRunner {
     try {
       sandbox.createFileSystem();
       OutErr outErr = policy.getFileOutErr();
+      policy.prefetchInputs();
+
       SpawnResult result = run(sandbox, outErr, timeout);
-  
+
       policy.lockOutputFiles();
-  
       try {
         // We copy the outputs even when the command failed.
         sandbox.copyOutputs(execRoot);