Refactor the sandbox spawn runners

- Move statisticsPath into SandboxedSpawn
- Unify all implementations of actuallyExec, inline into exec

This requires some changes to DockerSandboxedSpawnRunner that I'm not
sure about.

This is in preparation for merging local and sandboxed execution into a
single code path, in order to make it easier to support async execution.

PiperOrigin-RevId: 266748829

src/main/java/com/google/devtools/build/lib/sandbox/AbstractContainerizingSandboxedSpawn.java

@@ -30,6 +30,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
+import javax.annotation.Nullable;
 
 /**
  * Implements the general flow of a sandboxed spawn that uses a container directory to build an
@@ -50,6 +51,7 @@ public abstract class AbstractContainerizingSandboxedSpawn implements SandboxedS
   private final SandboxOutputs outputs;
   private final Set<Path> writableDirs;
   private final TreeDeleter treeDeleter;
+  private final Path statisticsPath;
 
   public AbstractContainerizingSandboxedSpawn(
       Path sandboxPath,
@@ -59,7 +61,8 @@ public AbstractContainerizingSandboxedSpawn(
       SandboxInputs inputs,
       SandboxOutputs outputs,
       Set<Path> writableDirs,
-      TreeDeleter treeDeleter) {
+      TreeDeleter treeDeleter,
+      @Nullable Path statisticsPath) {
     this.sandboxPath = sandboxPath;
     this.sandboxExecRoot = sandboxExecRoot;
     this.arguments = arguments;
@@ -68,6 +71,7 @@ public AbstractContainerizingSandboxedSpawn(
     this.outputs = outputs;
     this.writableDirs = writableDirs;
     this.treeDeleter = treeDeleter;
+    this.statisticsPath = statisticsPath;
   }
 
   @Override
@@ -85,6 +89,12 @@ public Map<String, String> getEnvironment() {
     return environment;
   }
 
+  @Override
+  @Nullable
+  public Path getStatisticsPath() {
+    return statisticsPath;
+  }
+
   @Override
   public void createFileSystem() throws IOException {
     createDirectories();

src/main/java/com/google/devtools/build/lib/sandbox/AbstractSandboxSpawnRunner.java

@@ -57,6 +57,7 @@ abstract class AbstractSandboxSpawnRunner implements SpawnRunner {
   private final boolean verboseFailures;
   private final ImmutableSet<Path> inaccessiblePaths;
   protected final BinTools binTools;
+  private final Path execRoot;
   private final ResourceManager resourceManager;
 
   public AbstractSandboxSpawnRunner(CommandEnvironment cmdEnv) {
@@ -65,18 +66,20 @@ public AbstractSandboxSpawnRunner(CommandEnvironment cmdEnv) {
     this.inaccessiblePaths =
         sandboxOptions.getInaccessiblePaths(cmdEnv.getRuntime().getFileSystem());
     this.binTools = cmdEnv.getBlazeWorkspace().getBinTools();
+    this.execRoot = cmdEnv.getExecRoot();
     this.resourceManager = cmdEnv.getLocalResourceManager();
   }
 
   @Override
-  public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
+  public final SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
       throws ExecException, IOException, InterruptedException {
     ActionExecutionMetadata owner = spawn.getResourceOwner();
     context.report(ProgressStatus.SCHEDULING, getName());
     try (ResourceHandle ignored =
         resourceManager.acquireResources(owner, spawn.getLocalResources())) {
       context.report(ProgressStatus.EXECUTING, getName());
-      return actuallyExec(spawn, context);
+      SandboxedSpawn sandbox = prepareSpawn(spawn, context);
+      return runSpawn(spawn, sandbox, context);
     } catch (IOException e) {
       throw new UserExecException("I/O exception during sandboxed execution", e);
     }
@@ -87,26 +90,18 @@ public boolean canExec(Spawn spawn) {
     return Spawns.mayBeSandboxed(spawn);
   }
 
-  // TODO(laszlocsomor): refactor this class to make `actuallyExec`'s contract clearer: the caller
-  // of `actuallyExec` should not depend on `actuallyExec` calling `runSpawn` because it's easy to
-  // forget to do so in `actuallyExec`'s implementations.
-  protected abstract SpawnResult actuallyExec(Spawn spawn, SpawnExecutionContext context)
-      throws ExecException, InterruptedException, IOException;
+  protected abstract SandboxedSpawn prepareSpawn(Spawn spawn, SpawnExecutionContext context)
+      throws IOException, ExecException;
 
-  protected SpawnResult runSpawn(
-      Spawn originalSpawn,
-      SandboxedSpawn sandbox,
-      SpawnExecutionContext context,
-      Path execRoot,
-      Duration timeout,
-      Path statisticsPath)
+  private SpawnResult runSpawn(
+      Spawn originalSpawn, SandboxedSpawn sandbox, SpawnExecutionContext context)
       throws IOException, InterruptedException {
     try {
       sandbox.createFileSystem();
       FileOutErr outErr = context.getFileOutErr();
       context.prefetchInputs();
 
-      SpawnResult result = run(originalSpawn, sandbox, outErr, timeout, statisticsPath);
+      SpawnResult result = run(originalSpawn, sandbox, context.getTimeout(), outErr);
 
       context.lockOutputFiles();
       try {
@@ -124,11 +119,7 @@ protected SpawnResult runSpawn(
   }
 
   private final SpawnResult run(
-      Spawn originalSpawn,
-      SandboxedSpawn sandbox,
-      FileOutErr outErr,
-      Duration timeout,
-      Path statisticsPath)
+      Spawn originalSpawn, SandboxedSpawn sandbox, Duration timeout, FileOutErr outErr)
       throws IOException, InterruptedException {
     String failureMessage;
     if (sandboxOptions.sandboxDebug) {
@@ -156,6 +147,10 @@ private final SpawnResult run(
     subprocessBuilder.setStderr(outErr.getErrorPath().getPathFile());
     subprocessBuilder.setEnv(sandbox.getEnvironment());
     subprocessBuilder.setArgv(sandbox.getArguments());
+    boolean useSubprocessTimeout = sandbox.useSubprocessTimeout();
+    if (useSubprocessTimeout) {
+      subprocessBuilder.setTimeoutMillis(timeout.toMillis());
+    }
     long startTime = System.currentTimeMillis();
     TerminationStatus terminationStatus;
     try {
@@ -182,9 +177,11 @@ private final SpawnResult run(
           .build();
     }
 
-    // TODO(b/62588075): Calculate wall time inside commands instead?
+    // TODO(b/62588075): Calculate wall time inside Subprocess instead?
     Duration wallTime = Duration.ofMillis(System.currentTimeMillis() - startTime);
-    boolean wasTimeout = terminationStatus.timedOut() || wasTimeout(timeout, wallTime);
+    boolean wasTimeout =
+        (useSubprocessTimeout && terminationStatus.timedOut())
+            || (!useSubprocessTimeout && wasTimeout(timeout, wallTime));
     int exitCode = wasTimeout ? POSIX_TIMEOUT_EXIT_CODE : terminationStatus.getRawExitCode();
     Status status =
         wasTimeout
@@ -199,6 +196,7 @@ private final SpawnResult run(
             .setWallTime(wallTime)
             .setFailureMessage(status != Status.SUCCESS || exitCode != 0 ? failureMessage : "");
 
+    Path statisticsPath = sandbox.getStatisticsPath();
     if (statisticsPath != null) {
       ExecutionStatistics.getResourceUsage(statisticsPath)
           .ifPresent(

src/main/java/com/google/devtools/build/lib/sandbox/CopyingSandboxedSpawn.java

@@ -25,11 +25,13 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import javax.annotation.Nullable;
 
 /**
  * Creates an execRoot for a Spawn that contains input files as copies from their original source.
  */
 public class CopyingSandboxedSpawn extends AbstractContainerizingSandboxedSpawn {
+  private final Runnable successCallback;
 
   public CopyingSandboxedSpawn(
       Path sandboxPath,
@@ -39,7 +41,9 @@ public CopyingSandboxedSpawn(
       SandboxInputs inputs,
       SandboxOutputs outputs,
       Set<Path> writableDirs,
-      TreeDeleter treeDeleter) {
+      TreeDeleter treeDeleter,
+      @Nullable Path statisticsPath,
+      Runnable successCallback) {
     super(
         sandboxPath,
         sandboxExecRoot,
@@ -48,7 +52,15 @@ public CopyingSandboxedSpawn(
         inputs,
         outputs,
         writableDirs,
-        treeDeleter);
+        treeDeleter,
+        statisticsPath);
+    this.successCallback = successCallback;
+  }
+
+  @Override
+  public void copyOutputs(Path execRoot) throws IOException {
+    successCallback.run();
+    super.copyOutputs(execRoot);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/sandbox/DarwinSandboxedSpawnRunner.java

@@ -21,10 +21,10 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.ByteStreams;
+import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionStrategy;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
-import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.Spawns;
 import com.google.devtools.build.lib.exec.TreeDeleter;
 import com.google.devtools.build.lib.exec.local.LocalEnvProvider;
@@ -213,8 +213,8 @@ private static String getConfStr(String confVar) throws IOException {
   }
 
   @Override
-  protected SpawnResult actuallyExec(Spawn spawn, SpawnExecutionContext context)
-      throws IOException, InterruptedException {
+  protected SandboxedSpawn prepareSpawn(Spawn spawn, SpawnExecutionContext context)
+      throws IOException, ExecException {
     // Each invocation of "exec" gets its own sandbox base.
     // Note that the value returned by context.getId() is only unique inside one given SpawnRunner,
     // so we have to prefix our name to turn it into a globally unique value.
@@ -274,54 +274,52 @@ protected SpawnResult actuallyExec(Spawn spawn, SpawnExecutionContext context)
             execRoot,
             getSandboxOptions().symlinkedSandboxExpandsTreeArtifactsInRunfilesTree);
 
-    SandboxedSpawn sandbox;
     if (sandboxfsProcess != null) {
-      sandbox =
-          new SandboxfsSandboxedSpawn(
-              sandboxfsProcess,
-              sandboxPath,
-              commandLine,
-              environment,
-              inputs,
-              outputs,
-              ImmutableSet.of(),
-              sandboxfsMapSymlinkTargets,
-              treeDeleter) {
-            @Override
-            public void createFileSystem() throws IOException {
-              super.createFileSystem();
-              writeConfig(
-                  sandboxConfigPath,
-                  writableDirs,
-                  getInaccessiblePaths(),
-                  allowNetworkForThisSpawn,
-                  statisticsPath);
-            }
-          };
+      return new SandboxfsSandboxedSpawn(
+          sandboxfsProcess,
+          sandboxPath,
+          commandLine,
+          environment,
+          inputs,
+          outputs,
+          ImmutableSet.of(),
+          sandboxfsMapSymlinkTargets,
+          treeDeleter,
+          statisticsPath) {
+        @Override
+        public void createFileSystem() throws IOException {
+          super.createFileSystem();
+          writeConfig(
+              sandboxConfigPath,
+              writableDirs,
+              getInaccessiblePaths(),
+              allowNetworkForThisSpawn,
+              statisticsPath);
+        }
+      };
     } else {
-      sandbox =
-          new SymlinkedSandboxedSpawn(
-              sandboxPath,
-              sandboxExecRoot,
-              commandLine,
-              environment,
-              inputs,
-              outputs,
+      return new SymlinkedSandboxedSpawn(
+          sandboxPath,
+          sandboxExecRoot,
+          commandLine,
+          environment,
+          inputs,
+          outputs,
+          writableDirs,
+          treeDeleter,
+          statisticsPath) {
+        @Override
+        public void createFileSystem() throws IOException {
+          super.createFileSystem();
+          writeConfig(
+              sandboxConfigPath,
               writableDirs,
-              treeDeleter) {
-            @Override
-            public void createFileSystem() throws IOException {
-              super.createFileSystem();
-              writeConfig(
-                  sandboxConfigPath,
-                  writableDirs,
-                  getInaccessiblePaths(),
-                  allowNetworkForThisSpawn,
-                  statisticsPath);
-            }
-          };
+              getInaccessiblePaths(),
+              allowNetworkForThisSpawn,
+              statisticsPath);
+        }
+      };
     }
-    return runSpawn(spawn, sandbox, context, execRoot, timeout, statisticsPath);
   }
 
   private void writeConfig(