Switch TestRunnerAction to use continuations

This rewrites execute() in terms of beginExecution and the resulting
continuation(s), and removes the old implementation. It implicitly
switches all the existing tests to the new code paths.

There's a small risk that this introduces a subtle behavioral
difference that isn't caught by the existing tests, but I'd say we
have pretty good test coverage.

Progress on #6394.

PiperOrigin-RevId: 240097352
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionContext.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionContext.java
index e86e7ed..9ef921b 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionContext.java
@@ -121,18 +121,13 @@
   interface TestRunnerSpawn {
     ActionExecutionContext getActionExecutionContext();
 
-    /** Execute the test, and handle the test runner protocol. */
-    TestAttemptResult execute() throws InterruptedException, IOException, ExecException;
-
     /**
      * Begin the test attempt execution. This may block until the test attempt is complete and
      * return a completed result, or it may return a continuation with a non-null future
      * representing asynchronous execution.
      */
-    default TestAttemptContinuation beginExecution()
-        throws InterruptedException, IOException, ExecException {
-      return TestAttemptContinuation.of(execute());
-    }
+    TestAttemptContinuation beginExecution()
+        throws InterruptedException, IOException, ExecException;
 
     /**
      * After the first attempt has run, this method is called to determine the maximum number of
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
index 2bf8670..b4ec9ea 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java
@@ -754,6 +754,13 @@
       throws InterruptedException, ActionExecutionException {
     TestActionContext testActionContext =
         actionExecutionContext.getContext(TestActionContext.class);
+    return beginExecution(actionExecutionContext, testActionContext);
+  }
+
+  @VisibleForTesting
+  public ActionContinuationOrResult beginExecution(
+      ActionExecutionContext actionExecutionContext, TestActionContext testActionContext)
+      throws InterruptedException, ActionExecutionException {
     try {
       TestRunnerSpawn testRunnerSpawn =
           testActionContext.createTestRunnerSpawn(this, actionExecutionContext);
@@ -792,79 +799,17 @@
       ActionExecutionContext actionExecutionContext, TestActionContext testActionContext)
       throws ActionExecutionException, InterruptedException {
     try {
-      TestRunnerSpawn testRunnerSpawn =
-          testActionContext.createTestRunnerSpawn(this, actionExecutionContext);
-      return ActionResult.create(
-          runAttempts(
-              testRunnerSpawn, actionExecutionContext, testActionContext.isTestKeepGoing()));
-    } catch (ExecException e) {
-      throw e.toActionExecutionException(this);
+      ActionContinuationOrResult continuation =
+          beginExecution(actionExecutionContext, testActionContext);
+      while (!continuation.isDone()) {
+        continuation = continuation.execute();
+      }
+      return continuation.get();
     } finally {
       unconditionalExecution = null;
     }
   }
 
-  private List<SpawnResult> runAttempts(
-      TestRunnerSpawn testRunnerSpawn,
-      ActionExecutionContext actionExecutionContext,
-      boolean keepGoing)
-      throws ExecException, InterruptedException {
-    List<SpawnResult> spawnResults = new ArrayList<>();
-    runAttempts(
-        testRunnerSpawn, actionExecutionContext, keepGoing, spawnResults, new ArrayList<>());
-    return spawnResults;
-  }
-
-  private void runAttempts(
-      TestRunnerSpawn testRunnerSpawn,
-      ActionExecutionContext actionExecutionContext,
-      boolean keepGoing,
-      List<SpawnResult> spawnResults,
-      List<FailedAttemptResult> failedAttempts)
-      throws ExecException, InterruptedException {
-    try {
-      TestAttemptResult testProcessResult = testRunnerSpawn.execute();
-      spawnResults.addAll(testProcessResult.spawnResults());
-      int maxAttempts = testRunnerSpawn.getMaxAttempts(testProcessResult);
-      // Failed test retry loop.
-      for (int attempt = 1; attempt < maxAttempts && !testProcessResult.hasPassed(); attempt++) {
-        failedAttempts.add(
-            testRunnerSpawn.finalizeFailedTestAttempt(
-                testProcessResult, failedAttempts.size() + 1));
-
-        testProcessResult = testRunnerSpawn.execute();
-        spawnResults.addAll(testProcessResult.spawnResults());
-      }
-
-      TestRunnerSpawn fallbackRunner;
-      if (!testProcessResult.hasPassed()
-          && (fallbackRunner = testRunnerSpawn.getFallbackRunner()) != null) {
-        // All attempts failed and fallback is enabled.
-        failedAttempts.add(
-            testRunnerSpawn.finalizeFailedTestAttempt(
-                testProcessResult, failedAttempts.size() + 1));
-        runAttempts(
-            fallbackRunner, actionExecutionContext, keepGoing, spawnResults, failedAttempts);
-      } else {
-        testRunnerSpawn.finalizeTest(testProcessResult, failedAttempts);
-
-        if (!keepGoing && !testProcessResult.hasPassed()) {
-          throw new TestExecException("Test failed: aborting");
-        }
-      }
-    } catch (IOException e) {
-      // Print the stack trace, otherwise the unexpected I/O error is hard to diagnose.
-      // A stack trace could help with bugs like https://github.com/bazelbuild/bazel/issues/4924
-      StringBuilder sb = new StringBuilder();
-      sb.append("Caught I/O exception: ").append(e.getMessage());
-      for (Object s : e.getStackTrace()) {
-        sb.append("\n\t").append(s);
-      }
-      actionExecutionContext.getEventHandler().handle(Event.error(sb.toString()));
-      throw new EnvironmentalExecException("unexpected I/O exception", e);
-    }
-  }
-
   @Override
   public String getMnemonic() {
     return MNEMONIC;
diff --git a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
index cc48367..39d2671 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
@@ -272,127 +272,6 @@
         .execute();
   }
 
-  private StandaloneTestResult executeTestAttempt(
-      TestRunnerAction action,
-      Spawn spawn,
-      ActionExecutionContext actionExecutionContext,
-      Path execRoot)
-      throws ExecException, IOException, InterruptedException {
-    Closeable streamed = null;
-    // We have two protos to represent test attempts:
-    // 1. com.google.devtools.build.lib.view.test.TestStatus.TestResultData represents both failed
-    //    attempts and finished tests. Bazel stores this to disk to persist cached test result
-    //    information across server restarts.
-    // 2. com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.TestResult
-    //    represents only individual attempts (failed or not). Bazel reports this as an event to the
-    //    Build Event Protocol, but never saves it to disk.
-    //
-    // The TestResult proto is always constructed from a TestResultData instance, either one that is
-    // created right here, or one that is read back from disk.
-    TestResultData.Builder builder = TestResultData.newBuilder();
-
-    SpawnActionContext spawnActionContext =
-        actionExecutionContext.getContext(SpawnActionContext.class);
-    List<SpawnResult> spawnResults = new ArrayList<>();
-
-    Path out = actionExecutionContext.getInputPath(action.getTestLog());
-    Path err = action.resolve(execRoot).getTestStderr();
-    long startTime = actionExecutionContext.getClock().currentTimeMillis();
-    try (FileOutErr testOutErr = new FileOutErr(out, err)) {
-      if (executionOptions.testOutput.equals(TestOutputFormat.STREAMED)) {
-        streamed =
-            new StreamedTestOutput(
-                Reporter.outErrForReporter(actionExecutionContext.getEventHandler()), out);
-      }
-      try {
-        spawnResults.addAll(
-            spawnActionContext.exec(spawn, actionExecutionContext.withFileOutErr(testOutErr)));
-        builder
-            .setTestPassed(true)
-            .setStatus(BlazeTestStatus.PASSED)
-            .setPassedLog(out.getPathString());
-      } catch (SpawnExecException e) {
-        // If this method returns normally, then the higher level will rerun the test (up to
-        // --flaky_test_attempts times).
-        if (e.isCatastrophic()) {
-          // Rethrow as the error was catastrophic and thus the build has to be halted.
-          throw e;
-        }
-        if (!e.getSpawnResult().setupSuccess()) {
-          // Rethrow as the test could not be run and thus there's no point in retrying.
-          throw e;
-        }
-        builder
-            .setTestPassed(false)
-            .setStatus(e.hasTimedOut() ? BlazeTestStatus.TIMEOUT : BlazeTestStatus.FAILED)
-            .addFailedLogs(out.getPathString());
-        spawnResults.add(e.getSpawnResult());
-      }
-      if (!testOutErr.hasRecordedOutput()) {
-        // Make sure that the test.log exists.
-        FileSystemUtils.touchFile(out);
-      }
-      // Append any error output to the test.log. This is very rare.
-      appendStderr(testOutErr);
-    }
-
-    long endTime = actionExecutionContext.getClock().currentTimeMillis();
-    long duration = endTime - startTime;
-    // SpawnActionContext guarantees the first entry to correspond to the spawn passed in (there may
-    // be additional entries due to tree artifact handling).
-    SpawnResult primaryResult = spawnResults.get(0);
-
-    // The SpawnResult of a remotely cached or remotely executed action may not have walltime
-    // set. We fall back to the time measured here for backwards compatibility.
-    duration = primaryResult.getWallTime().orElse(Duration.ofMillis(duration)).toMillis();
-    BuildEventStreamProtos.TestResult.ExecutionInfo.Builder executionInfo =
-        extractExecutionInfo(primaryResult, builder);
-
-    builder.setStartTimeMillisEpoch(startTime);
-    builder.addTestTimes(duration);
-    builder.addTestProcessTimes(duration);
-    builder.setRunDurationMillis(duration);
-    if (streamed != null) {
-      streamed.close();
-    }
-
-    // If the test did not create a test.xml, and --experimental_split_xml_generation is enabled,
-    // then we run a separate action to create a test.xml from test.log. We do this as a spawn
-    // rather than doing it locally in-process, as the test.log file may only exist remotely (when
-    // remote execution is enabled), and we do not want to have to download it.
-    Path xmlOutputPath = action.resolve(actionExecutionContext.getExecRoot()).getXmlOutputPath();
-    if (executionOptions.splitXmlGeneration
-        && action.getTestLog().getPath().exists()
-        && !xmlOutputPath.exists()) {
-      Spawn xmlGeneratingSpawn = createXmlGeneratingSpawn(action, primaryResult);
-      // We treat all failures to generate the test.xml here as catastrophic, and won't rerun
-      // the test if this fails. We redirect the output to a temporary file.
-      try (FileOutErr xmlSpawnOutErr = actionExecutionContext.getFileOutErr().childOutErr()) {
-        spawnResults.addAll(
-            spawnActionContext.exec(
-                xmlGeneratingSpawn, actionExecutionContext.withFileOutErr(xmlSpawnOutErr)));
-      }
-    }
-
-    TestCase details = parseTestResult(xmlOutputPath);
-    if (details != null) {
-      builder.setTestCase(details);
-    }
-
-    if (action.isCoverageMode()) {
-      builder.setHasCoverage(true);
-    }
-
-    return StandaloneTestResult.builder()
-        .setSpawnResults(spawnResults)
-        // We return the TestResultData.Builder rather than the finished TestResultData instance,
-        // as we may have to rename the output files in case the test needs to be rerun (if it
-        // failed here _and_ is marked flaky _and_ the number of flaky attempts is larger than 1).
-        .setTestResultDataBuilder(builder)
-        .setExecutionInfo(executionInfo.build())
-        .build();
-  }
-
   /** In rare cases, we might write something to stderr. Append it to the real test.log. */
   private static void appendStderr(FileOutErr outErr) throws IOException {
     Path stdErr = outErr.getErrorPath();
@@ -618,12 +497,6 @@
     }
 
     @Override
-    public TestAttemptResult execute() throws InterruptedException, IOException, ExecException {
-      prepareFileSystem(testAction, tmpDir, coverageDir, workingDirectory);
-      return executeTestAttempt(testAction, spawn, actionExecutionContext, execRoot);
-    }
-
-    @Override
     public int getMaxAttempts(TestAttemptResult firstTestAttemptResult) {
       return getTestAttempts(testAction);
     }
diff --git a/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java b/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
index 219d89e..1ebc38c 100644
--- a/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
@@ -32,6 +32,7 @@
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.SpawnResult.Status;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
@@ -188,7 +189,8 @@
             .setWallTime(Duration.ofMillis(10))
             .setRunnerName("test")
             .build();
-    when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
+    when(spawnActionContext.beginExecution(any(), any()))
+        .thenReturn(SpawnContinuation.immediate(expectedSpawnResult));
 
     ActionExecutionContext actionExecutionContext =
         new FakeActionExecutionContext(createTempOutErr(tmpDirRoot), spawnActionContext);
@@ -254,9 +256,9 @@
             .setWallTime(Duration.ofMillis(15))
             .setRunnerName("test")
             .build();
-    when(spawnActionContext.exec(any(), any()))
+    when(spawnActionContext.beginExecution(any(), any()))
         .thenThrow(new SpawnExecException("test failed", failSpawnResult, false))
-        .thenReturn(ImmutableList.of(passSpawnResult));
+        .thenReturn(SpawnContinuation.immediate(passSpawnResult));
 
     ActionExecutionContext actionExecutionContext =
         new FakeActionExecutionContext(createTempOutErr(tmpDirRoot), spawnActionContext);
@@ -322,7 +324,8 @@
             .setRunnerName("remote")
             .setExecutorHostname("a-remote-host")
             .build();
-    when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
+    when(spawnActionContext.beginExecution(any(), any()))
+        .thenReturn(SpawnContinuation.immediate(expectedSpawnResult));
 
     ActionExecutionContext actionExecutionContext =
         new FakeActionExecutionContext(createTempOutErr(tmpDirRoot), spawnActionContext);
@@ -380,7 +383,8 @@
             .setWallTime(Duration.ofMillis(10))
             .setRunnerName("remote cache")
             .build();
-    when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
+    when(spawnActionContext.beginExecution(any(), any()))
+        .thenReturn(SpawnContinuation.immediate(expectedSpawnResult));
 
     ActionExecutionContext actionExecutionContext =
         new FakeActionExecutionContext(createTempOutErr(tmpDirRoot), spawnActionContext);
@@ -438,7 +442,7 @@
             .setExitCode(1)
             .setRunnerName("test")
             .build();
-    when(spawnActionContext.exec(any(), any()))
+    when(spawnActionContext.beginExecution(any(), any()))
         .thenAnswer(
             (invocation) -> {
               Spawn spawn = invocation.getArgument(0);
@@ -456,7 +460,7 @@
                     /*forciblyRunRemotely=*/ false,
                     /*catastrophe=*/ false);
               } else {
-                return ImmutableList.of(
+                return SpawnContinuation.immediate(
                     new SpawnResult.Builder()
                         .setStatus(Status.SUCCESS)
                         .setRunnerName("test")
@@ -535,7 +539,7 @@
             .setRunnerName("test")
             .build();
     List<FileOutErr> called = new ArrayList<>();
-    when(spawnActionContext.exec(any(), any()))
+    when(spawnActionContext.beginExecution(any(), any()))
         .thenAnswer(
             (invocation) -> {
               Spawn spawn = invocation.getArgument(0);
@@ -561,7 +565,7 @@
                         ? "standalone/failing_test.exe"
                         : "standalone/failing_test";
                 assertThat(spawn.getEnvironment()).containsEntry("TEST_BINARY", testName);
-                return ImmutableList.of(xmlGeneratorSpawnResult);
+                return SpawnContinuation.immediate(xmlGeneratorSpawnResult);
               }
             });
 
@@ -618,7 +622,8 @@
 
     SpawnResult expectedSpawnResult =
         new SpawnResult.Builder().setStatus(Status.SUCCESS).setRunnerName("test").build();
-    when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
+    when(spawnActionContext.beginExecution(any(), any()))
+        .thenReturn(SpawnContinuation.immediate(expectedSpawnResult));
 
     FileOutErr outErr = createTempOutErr(tmpDirRoot);
     ActionExecutionContext actionExecutionContext =