Bring back the persistent test runner functionality, on ExperimentalTestRunner.

This is essentially a rollforward of commit 7d0561b6ca92d72bd8767d4dca50e5437976812c, and changes triggering the perisitent runner using an environment variable instead of argument as suggested in commit 7d0561b6ca92d72bd8767d4dca50e5437976812c

--
PiperOrigin-RevId: 149540564
MOS_MIGRATED_REVID=149540564

4 files changed, 75 insertions, 11 deletions

diff --git a/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BUILD b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BUILD
index c5a8006794..464e838630 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BUILD
+++ b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BUILD
@@ -49,6 +49,7 @@ java_library(
         "//src/java_tools/junitrunner/java/com/google/testing/junit/runner/model",
         "//src/java_tools/junitrunner/java/com/google/testing/junit/runner/sharding",
         "//src/java_tools/junitrunner/java/com/google/testing/junit/runner/util",
+        "//src/main/protobuf:worker_protocol_java_proto",
         "//third_party:junit4",
     ],
 )
diff --git a/src/java_tools/junitrunner/java/com/google/testing/junit/runner/ExperimentalTestRunner.java b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/ExperimentalTestRunner.java
index f13915e8e9..e5455259be 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/junit/runner/ExperimentalTestRunner.java
+++ b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/ExperimentalTestRunner.java
@@ -14,10 +14,14 @@
 
 package com.google.testing.junit.runner;
 
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
 import com.google.testing.junit.runner.internal.StackTraces;
 import com.google.testing.junit.runner.junit4.JUnit4InstanceModules.Config;
 import com.google.testing.junit.runner.junit4.JUnit4InstanceModules.SuiteClass;
 import com.google.testing.junit.runner.junit4.JUnit4Runner;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.PrintStream;
 import java.text.DateFormat;
 import java.text.SimpleDateFormat;
@@ -58,13 +62,13 @@ public class ExperimentalTestRunner {
     PrintStream stderr = System.err;
 
     String suiteClassName = System.getProperty(TEST_SUITE_PROPERTY_NAME);
-    System.out.println("WARNING: RUNNING EXPERIMENTAL TEST RUNNER");
 
-    if (args.length >= 1 && args[args.length - 1].equals("--persistent_test_runner")) {
-      System.err.println("Requested test strategy is currently unsupported.");
-      System.exit(1);
+    if ("true".equals(System.getenv("PERSISTENT_TEST_RUNNER"))) {
+      System.exit(runPersistentTestRunner(suiteClassName));
     }
 
+    System.out.println("WARNING: RUNNING EXPERIMENTAL TEST RUNNER");
+
     if (!checkTestSuiteProperty(suiteClassName)) {
       System.exit(2);
     }
@@ -134,6 +138,44 @@ public class ExperimentalTestRunner {
     return runner.run().wasSuccessful() ? 0 : 1;
   }
 
+  private static int runPersistentTestRunner(String suiteClassName) {
+    PrintStream originalStdOut = System.out;
+    PrintStream originalStdErr = System.err;
+
+    while (true) {
+      try {
+        WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);
+
+        if (request == null) {
+          break;
+        }
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        PrintStream ps = new PrintStream(baos, true);
+        System.setOut(ps);
+        System.setErr(ps);
+
+        String[] arguments = request.getArgumentsList().toArray(new String[0]);
+        int exitCode = -1;
+        try {
+          exitCode = runTestsInSuite(suiteClassName, arguments);
+        } finally {
+          System.setOut(originalStdOut);
+          System.setErr(originalStdErr);
+        }
+
+        WorkResponse response =
+            WorkResponse.newBuilder().setOutput(baos.toString()).setExitCode(exitCode).build();
+        response.writeDelimitedTo(System.out);
+        System.out.flush();
+
+      } catch (IOException e) {
+        e.printStackTrace();
+        return 1;
+      }
+    }
+    return 0;
+  }
+
   private static Class<?> getTestClass(String name) {
     if (name == null) {
       return null;
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerTestStrategy.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerTestStrategy.java
index 36ab7bf8e5..a3b8dff3d3 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerTestStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerTestStrategy.java
@@ -87,7 +87,7 @@ public class WorkerTestStrategy extends StandaloneTestStrategy {
       throw new UserExecException("Tests that do not use the default test runner are incompatible"
           + " with the persistent worker test strategy. Please use another test strategy");
     }
-    if (!action.isCoverageMode()) {
+    if (action.isCoverageMode()) {
       throw new UserExecException("Coverage is currently incompatible"
           + " with the persistent worker test strategy. Please use another test strategy");
     }
@@ -96,7 +96,7 @@ public class WorkerTestStrategy extends StandaloneTestStrategy {
     return execInWorker(
         action,
         actionExecutionContext,
-        spawn.getEnvironment(),
+        addPersistentRunnerVars(spawn.getEnvironment()),
         startupArgs,
         actionExecutionContext.getExecutor().getExecRoot(),
         maxRetries);
@@ -111,6 +111,14 @@ public class WorkerTestStrategy extends StandaloneTestStrategy {
       int retriesLeft)
       throws ExecException, InterruptedException, IOException {
     Executor executor = actionExecutionContext.getExecutor();
+
+    // TODO(kush): Remove once we're out of the experimental phase.
+    executor
+        .getEventHandler()
+        .handle(
+            Event.warn(
+                "RUNNING TEST IN AN EXPERIMENTAL PERSISTENT WORKER. RESULTS MAY BE INACCURATE"));
+
     TestResultData.Builder builder = TestResultData.newBuilder();
 
     Path testLogPath = action.getTestLog().getPath();
@@ -195,6 +203,19 @@ public class WorkerTestStrategy extends StandaloneTestStrategy {
     }
   }
 
+  private static Map<String, String> addPersistentRunnerVars(Map<String, String> originalEnv)
+      throws UserExecException {
+    if (originalEnv.containsKey("PERSISTENT_TEST_RUNNER")) {
+      throw new UserExecException(
+          "Found clashing environment variable with persistent_test_runner."
+              + " Please use another test strategy");
+    }
+    return ImmutableMap.<String, String>builder()
+        .putAll(originalEnv)
+        .put("PERSISTENT_TEST_RUNNER", "true")
+        .build();
+  }
+
   private List<String> getStartUpArgs(TestRunnerAction action) throws ExecException {
     List<String> args = getArgs(/*coverageScript=*/ "coverage-is-not-supported", action);
     ImmutableList.Builder<String> startupArgs = ImmutableList.builder();
@@ -202,8 +223,6 @@ public class WorkerTestStrategy extends StandaloneTestStrategy {
     startupArgs.add(args.get(0)).add("--no_echo");
     // Add remaining of the original args.
     startupArgs.addAll(args.subList(1, args.size()));
-    // Make the Test runner run persistently.
-    startupArgs.add("--persistent_test_runner");
     // Add additional flags requested for this invocation.
     startupArgs.addAll(MoreObjects.firstNonNull(
             extraFlags.get(action.getMnemonic()), ImmutableList.<String>of()));
diff --git a/src/test/shell/bazel/persistent_test_runner_test.sh b/src/test/shell/bazel/persistent_test_runner_test.sh
index 8af6d04a65..db7a15b70a 100755
--- a/src/test/shell/bazel/persistent_test_runner_test.sh
+++ b/src/test/shell/bazel/persistent_test_runner_test.sh
@@ -22,7 +22,7 @@ CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 source "${CURRENT_DIR}/../integration_test_setup.sh" \
   || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
 
-function DISABLED_test_simple_scenario() {
+function test_simple_scenario() {
   mkdir -p java/testrunners || fail "mkdir failed"
 
   cat > java/testrunners/TestsPass.java <<EOF
@@ -63,12 +63,14 @@ EOF
   cat > java/testrunners/BUILD <<EOF
 java_test(name = "TestsPass",
           srcs = ['TestsPass.java'],
-          deps = ['@bazel_tools//tools/jdk:TestRunner_deploy.jar'],
+          tags = ["experimental_testrunner"],
+          deps = ['@bazel_tools//tools/jdk:ExperimentalTestRunner_deploy.jar'],
 )
 
 java_test(name = "TestsFail",
           srcs = ['TestsFail.java'],
-          deps = ['@bazel_tools//tools/jdk:TestRunner_deploy.jar'],
+          tags = ["experimental_testrunner"],
+          deps = ['@bazel_tools//tools/jdk:ExperimentalTestRunner_deploy.jar'],
 )
 EOF