Extend the SpawnRunner API

- add an id for logging; this allows us to correlate log entries for the same
  spawn from multiple spawn runner implementations in the future
- add a prefetch method to the SpawnExecutionPolicy; better than relying on
  the ActionInputPrefetcher being injected in the constructor
- add a name parameter to the report method; this is in preparation for a
  single unified SpawnStrategy implementation - it's basically the last bit of
  difference between SandboxStrategy and RemoteSpawnStrategy; they're otherwise
  equivalent (if not identical)

PiperOrigin-RevId: 162194684

8 files changed, 105 insertions, 46 deletions

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 9445171656..c2819766cc 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
@@ -120,7 +120,38 @@ public interface SpawnRunner {
    * be used by different threads, so they MUST not call any shared non-thread-safe objects.
    */
   public interface SpawnExecutionPolicy {
-    /** The input file cache for this specific spawn. */
+    /**
+     * Returns a unique id for this spawn, to be used for logging. Note that a single spawn may be
+     * passed to multiple {@link SpawnRunner} implementations, so any log entries should also
+     * contain the identity of the spawn runner implementation.
+     */
+    int getId();
+
+    /**
+     * Prefetches the Spawns input files to the local machine. There are cases where Bazel runs on a
+     * network file system, and prefetching the files in parallel is a significant performance win.
+     * This should only be called by local strategies when local execution is imminent.
+     *
+     * <p>Should be called with the equivalent of:
+     * <code>
+     * policy.prefetchInputs(
+     *      Iterables.filter(policy.getInputMapping().values(), Predicates.notNull()));
+     * </code>
+     *
+     * <p>Note in particular that {@link #getInputMapping} may return {@code null} values, but
+     * this method does not accept {@code null} values.
+     *
+     * <p>The reason why this method requires passing in the inputs is that getInputMapping may be
+     * slow to compute, so if the implementation already called it, we don't want to compute it
+     * 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;
+
+    /**
+     * The input file metadata cache for this specific spawn, which can be used to efficiently
+     * obtain file digests and sizes.
+     */
     ActionInputFileCache getActionInputFileCache();
 
     /** An artifact expander. */
@@ -146,7 +177,7 @@ public interface SpawnRunner {
     SortedMap<PathFragment, ActionInput> getInputMapping() throws IOException;
 
     /** Reports a progress update to the Spawn strategy. */
-    void report(ProgressStatus state);
+    void report(ProgressStatus state, String name);
   }
 
   /**
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 af968ae114..2619ef5249 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
@@ -27,7 +27,6 @@ import com.google.devtools.build.lib.actions.ResourceManager.ResourceHandle;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.Spawns;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
-import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.SpawnResult;
 import com.google.devtools.build.lib.exec.SpawnResult.Status;
 import com.google.devtools.build.lib.exec.SpawnRunner;
@@ -47,7 +46,6 @@ import java.util.ArrayList;
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -69,9 +67,6 @@ public final class LocalSpawnRunner implements SpawnRunner {
   private final ResourceManager resourceManager;
 
   private final String hostName;
-  private final AtomicInteger execCount;
-
-  private final ActionInputPrefetcher actionInputPrefetcher;
 
   private final LocalExecutionOptions localExecutionOptions;
 
@@ -86,9 +81,7 @@ public final class LocalSpawnRunner implements SpawnRunner {
   }
 
   public LocalSpawnRunner(
-      AtomicInteger execCount,
       Path execRoot,
-      ActionInputPrefetcher actionInputPrefetcher,
       LocalExecutionOptions localExecutionOptions,
       ResourceManager resourceManager,
       boolean useProcessWrapper,
@@ -96,11 +89,9 @@ public final class LocalSpawnRunner implements SpawnRunner {
       String productName,
       LocalEnvProvider localEnvProvider) {
     this.execRoot = execRoot;
-    this.actionInputPrefetcher = Preconditions.checkNotNull(actionInputPrefetcher);
     this.processWrapper = getProcessWrapper(execRoot, localOs).getPathString();
     this.localExecutionOptions = Preconditions.checkNotNull(localExecutionOptions);
     this.hostName = NetUtil.findShortHostName();
-    this.execCount = execCount;
     this.resourceManager = resourceManager;
     this.useProcessWrapper = useProcessWrapper;
     this.productName = productName;
@@ -109,15 +100,12 @@ public final class LocalSpawnRunner implements SpawnRunner {
 
   public LocalSpawnRunner(
       Path execRoot,
-      ActionInputPrefetcher actionInputPrefetcher,
       LocalExecutionOptions localExecutionOptions,
       ResourceManager resourceManager,
       String productName,
       LocalEnvProvider localEnvProvider) {
     this(
-        new AtomicInteger(),
         execRoot,
-        actionInputPrefetcher,
         localExecutionOptions,
         resourceManager,
         OS.getCurrent() != OS.WINDOWS && getProcessWrapper(execRoot, OS.getCurrent()).exists(),
@@ -131,10 +119,10 @@ public final class LocalSpawnRunner implements SpawnRunner {
       Spawn spawn,
       SpawnExecutionPolicy policy) throws IOException, InterruptedException {
     ActionExecutionMetadata owner = spawn.getResourceOwner();
-    policy.report(ProgressStatus.SCHEDULING);
+    policy.report(ProgressStatus.SCHEDULING, "local");
     try (ResourceHandle handle =
         resourceManager.acquireResources(owner, spawn.getLocalResources())) {
-      policy.report(ProgressStatus.EXECUTING);
+      policy.report(ProgressStatus.EXECUTING, "local");
       policy.lockOutputFiles();
       return new SubprocessHandler(spawn, policy).run();
     }
@@ -149,7 +137,7 @@ public final class LocalSpawnRunner implements SpawnRunner {
     private State currentState = State.INITIALIZING;
     private final Map<State, Long> stateTimes = new EnumMap<>(State.class);
 
-    private final int id = execCount.getAndIncrement();
+    private final int id;
 
     public SubprocessHandler(
         Spawn spawn,
@@ -157,6 +145,7 @@ public final class LocalSpawnRunner implements SpawnRunner {
       Preconditions.checkArgument(!spawn.getArguments().isEmpty());
       this.spawn = spawn;
       this.policy = policy;
+      this.id = policy.getId();
       setState(State.PARSING);
     }
 
@@ -240,7 +229,7 @@ public final class LocalSpawnRunner implements SpawnRunner {
       if (Spawns.shouldPrefetchInputsForLocalExecution(spawn)) {
         stepLog(INFO, "prefetching inputs for local execution");
         setState(State.PREFETCHING_LOCAL_INPUTS);
-        actionInputPrefetcher.prefetchFiles(
+        policy.prefetchInputs(
             Iterables.filter(policy.getInputMapping().values(), Predicates.notNull()));
       }
 
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 b68cc993d8..c9e0655c71 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
@@ -60,17 +60,16 @@ final class RemoteActionContextProvider extends ActionContextProvider {
     spawnRunner = new RemoteSpawnRunner(
         env.getExecRoot(),
         remoteOptions,
-        createFallbackRunner(actionInputPrefetcher),
+        createFallbackRunner(),
         cache,
         executor);
     spawnStrategy =
         new RemoteSpawnStrategy(
-            "remote",
             spawnRunner,
             executionOptions.verboseFailures);
   }
 
-  private SpawnRunner createFallbackRunner(ActionInputPrefetcher actionInputPrefetcher) {
+  private SpawnRunner createFallbackRunner() {
     LocalExecutionOptions localExecutionOptions =
         env.getOptions().getOptions(LocalExecutionOptions.class);
     LocalEnvProvider localEnvProvider = OS.getCurrent() == OS.DARWIN
@@ -79,7 +78,6 @@ final class RemoteActionContextProvider extends ActionContextProvider {
     return
         new LocalSpawnRunner(
             env.getExecRoot(),
-            actionInputPrefetcher,
             localExecutionOptions,
             ResourceManager.instance(),
             env.getRuntime().getProductName(),
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java
index d0f95689d4..abe8fccde7 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java
@@ -83,7 +83,7 @@ final class RemoteSpawnRunner implements SpawnRunner {
       return fallbackRunner.exec(spawn, policy);
     }
 
-    policy.report(ProgressStatus.EXECUTING);
+    policy.report(ProgressStatus.EXECUTING, "remote");
     // Temporary hack: the TreeNodeRepository should be created and maintained upstream!
     ActionInputFileCache inputFileCache = policy.getActionInputFileCache();
     TreeNodeRepository repository = new TreeNodeRepository(execRoot, inputFileCache);
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
index 44d62b6493..7af46aebcd 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
@@ -27,6 +27,7 @@ import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.actions.Spawns;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.SpawnExecException;
 import com.google.devtools.build.lib.exec.SpawnInputExpander;
 import com.google.devtools.build.lib.exec.SpawnResult;
@@ -43,6 +44,7 @@ import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.util.SortedMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Strategy that uses a distributed cache for sharing action input and output files. Optionally this
@@ -54,14 +56,15 @@ import java.util.concurrent.TimeUnit;
 )
 final class RemoteSpawnStrategy implements SpawnActionContext {
   private final SpawnInputExpander spawnInputExpander = new SpawnInputExpander(/*strict=*/false);
-  private final String strategyName;
   private final SpawnRunner spawnRunner;
   private final boolean verboseFailures;
+  private final ActionInputPrefetcher inputPrefetcher;
+  private final AtomicInteger execCount = new AtomicInteger();
 
-  RemoteSpawnStrategy(String strategyName, SpawnRunner spawnRunner, boolean verboseFailures) {
-    this.strategyName = strategyName;
+  RemoteSpawnStrategy(SpawnRunner spawnRunner, boolean verboseFailures) {
     this.spawnRunner = spawnRunner;
     this.verboseFailures = verboseFailures;
+    this.inputPrefetcher = ActionInputPrefetcher.NONE;
   }
 
   @Override
@@ -83,7 +86,20 @@ final class RemoteSpawnStrategy implements SpawnActionContext {
     if (actionExecutionContext.reportsSubcommands()) {
       actionExecutionContext.reportSubcommand(spawn);
     }
+    final int timeoutSeconds = Spawns.getTimeoutSeconds(spawn);
     SpawnExecutionPolicy policy = new SpawnExecutionPolicy() {
+      private final int id = execCount.incrementAndGet();
+
+      @Override
+      public int getId() {
+        return id;
+      }
+
+      @Override
+      public void prefetchInputs(Iterable<ActionInput> inputs) throws IOException {
+        inputPrefetcher.prefetchFiles(inputs);
+      }
+
       @Override
       public ActionInputFileCache getActionInputFileCache() {
         return actionExecutionContext.getActionInputFileCache();
@@ -102,12 +118,7 @@ final class RemoteSpawnStrategy implements SpawnActionContext {
 
       @Override
       public long getTimeoutMillis() {
-        try {
-          return TimeUnit.SECONDS.toMillis(Spawns.getTimeoutSeconds(spawn));
-        } catch (ExecException e) {
-          // The exec info is set internally, so we can never fail to parse the timeout.
-          throw new RuntimeException(e);
-        }
+        return TimeUnit.SECONDS.toMillis(timeoutSeconds);
       }
 
       @Override
@@ -125,12 +136,12 @@ final class RemoteSpawnStrategy implements SpawnActionContext {
       }
 
       @Override
-      public void report(ProgressStatus state) {
+      public void report(ProgressStatus state, String name) {
         EventBus eventBus = actionExecutionContext.getEventBus();
         switch (state) {
           case EXECUTING:
             eventBus.post(
-                ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), strategyName));
+                ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), name));
             break;
           case SCHEDULING:
             eventBus.post(ActionStatusMessage.schedulingStrategy(spawn.getResourceOwner()));
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 d1b6e98452..8a4d9c4b62 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
@@ -66,10 +66,10 @@ abstract class AbstractSandboxSpawnRunner implements SpawnRunner {
   public SpawnResult exec(Spawn spawn, SpawnExecutionPolicy policy)
       throws ExecException, InterruptedException {
     ActionExecutionMetadata owner = spawn.getResourceOwner();
-    policy.report(ProgressStatus.SCHEDULING);
+    policy.report(ProgressStatus.SCHEDULING, getName());
     try (ResourceHandle ignored =
         ResourceManager.instance().acquireResources(owner, spawn.getLocalResources())) {
-      policy.report(ProgressStatus.EXECUTING);
+      policy.report(ProgressStatus.EXECUTING, getName());
       return actuallyExec(spawn, policy);
     } catch (IOException e) {
       throw new UserExecException("I/O exception during sandboxed execution", e);
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxStrategy.java b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxStrategy.java
index 7b325f61dc..bea49b7cbd 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxStrategy.java
@@ -25,6 +25,7 @@ import com.google.devtools.build.lib.actions.SandboxedSpawnActionContext;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.actions.Spawns;
+import com.google.devtools.build.lib.exec.ActionInputPrefetcher;
 import com.google.devtools.build.lib.exec.SpawnExecException;
 import com.google.devtools.build.lib.exec.SpawnInputExpander;
 import com.google.devtools.build.lib.exec.SpawnResult;
@@ -38,6 +39,7 @@ import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.util.SortedMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 /** Abstract common ancestor for sandbox strategies implementing the common parts. */
@@ -45,6 +47,8 @@ abstract class SandboxStrategy implements SandboxedSpawnActionContext {
   private final boolean verboseFailures;
   private final SpawnInputExpander spawnInputExpander;
   private final AbstractSandboxSpawnRunner spawnRunner;
+  private final ActionInputPrefetcher inputPrefetcher;
+  private final AtomicInteger execCount = new AtomicInteger();
 
   public SandboxStrategy(
       boolean verboseFailures,
@@ -52,6 +56,7 @@ abstract class SandboxStrategy implements SandboxedSpawnActionContext {
     this.verboseFailures = verboseFailures;
     this.spawnInputExpander = new SpawnInputExpander(false);
     this.spawnRunner = spawnRunner;
+    this.inputPrefetcher = ActionInputPrefetcher.NONE;
   }
 
   @Override
@@ -76,8 +81,19 @@ abstract class SandboxStrategy implements SandboxedSpawnActionContext {
       actionExecutionContext.reportSubcommand(spawn);
     }
     final int timeoutSeconds = Spawns.getTimeoutSeconds(spawn);
-    final EventBus eventBus = actionExecutionContext.getEventBus();
     SpawnExecutionPolicy policy = new SpawnExecutionPolicy() {
+      private final int id = execCount.incrementAndGet();
+
+      @Override
+      public int getId() {
+        return id;
+      }
+
+      @Override
+      public void prefetchInputs(Iterable<ActionInput> inputs) throws IOException {
+        inputPrefetcher.prefetchFiles(inputs);
+      }
+
       @Override
       public ActionInputFileCache getActionInputFileCache() {
         return actionExecutionContext.getActionInputFileCache();
@@ -118,13 +134,12 @@ abstract class SandboxStrategy implements SandboxedSpawnActionContext {
       }
 
       @Override
-      public void report(ProgressStatus state) {
+      public void report(ProgressStatus state, String name) {
         // TODO(ulfjack): We should report more details to the UI.
+        EventBus eventBus = actionExecutionContext.getEventBus();
         switch (state) {
           case EXECUTING:
-            String strategyName = spawnRunner.getName();
-            eventBus.post(
-                ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), strategyName));
+            eventBus.post(ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), name));
             break;
           case SCHEDULING:
             eventBus.post(ActionStatusMessage.schedulingStrategy(spawn.getResourceOwner()));
diff --git a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategy.java
index 4a637d0001..a32e99c9ae 100644
--- a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategy.java
@@ -45,6 +45,7 @@ import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.util.SortedMap;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Strategy that uses subprocessing to execute a process.
@@ -52,19 +53,21 @@ import java.util.SortedMap;
 @ExecutionStrategy(name = { "standalone", "local" }, contextType = SpawnActionContext.class)
 public class StandaloneSpawnStrategy implements SpawnActionContext {
   private final boolean verboseFailures;
+  private final ActionInputPrefetcher actionInputPrefetcher;
   private final LocalSpawnRunner localSpawnRunner;
+  private final AtomicInteger execCount = new AtomicInteger();
 
   public StandaloneSpawnStrategy(
       Path execRoot, ActionInputPrefetcher actionInputPrefetcher,
       LocalExecutionOptions localExecutionOptions, boolean verboseFailures, String productName,
       ResourceManager resourceManager) {
+    this.actionInputPrefetcher = actionInputPrefetcher;
     this.verboseFailures = verboseFailures;
     LocalEnvProvider localEnvProvider = OS.getCurrent() == OS.DARWIN
         ? new XCodeLocalEnvProvider()
         : LocalEnvProvider.UNMODIFIED;
     this.localSpawnRunner = new LocalSpawnRunner(
         execRoot,
-        actionInputPrefetcher,
         localExecutionOptions,
         resourceManager,
         productName,
@@ -80,6 +83,20 @@ public class StandaloneSpawnStrategy implements SpawnActionContext {
     final int timeoutSeconds = Spawns.getTimeoutSeconds(spawn);
     final EventBus eventBus = actionExecutionContext.getEventBus();
     SpawnExecutionPolicy policy = new SpawnExecutionPolicy() {
+      private final int id = execCount.incrementAndGet();
+
+      @Override
+      public int getId() {
+        return id;
+      }
+
+      @Override
+      public void prefetchInputs(Iterable<ActionInput> inputs) throws IOException {
+        if (Spawns.shouldPrefetchInputsForLocalExecution(spawn)) {
+          actionInputPrefetcher.prefetchFiles(inputs);
+        }
+      }
+
       @Override
       public ActionInputFileCache getActionInputFileCache() {
         return actionExecutionContext.getActionInputFileCache();
@@ -116,12 +133,10 @@ public class StandaloneSpawnStrategy implements SpawnActionContext {
       }
 
       @Override
-      public void report(ProgressStatus state) {
+      public void report(ProgressStatus state, String name) {
         switch (state) {
           case EXECUTING:
-            String strategyName = "local";
-            eventBus.post(
-                ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), strategyName));
+            eventBus.post(ActionStatusMessage.runningStrategy(spawn.getResourceOwner(), name));
             break;
           case SCHEDULING:
             eventBus.post(ActionStatusMessage.schedulingStrategy(spawn.getResourceOwner()));