Initial code for Persistent Java Test Runner.

At this point this does nothing more than re-run the exact same test without having to re-start the test runner.

In future iterations the aim is to be able to re-run tests with modified code, without having to re-start the test runner.

To test out the WorkerTestStrategy simply use --test_strategy=experimental_worker for a test with bazel.

--
PiperOrigin-RevId: 141465929
MOS_MIGRATED_REVID=141465929

11 files changed, 533 insertions, 52 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 1699cb882b..71a715a85a 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
@@ -17,6 +17,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/BazelTestRunner.java b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BazelTestRunner.java
index 6e9cf02188..7e3984138e 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BazelTestRunner.java
+++ b/src/java_tools/junitrunner/java/com/google/testing/junit/runner/BazelTestRunner.java
@@ -14,12 +14,16 @@
 
 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 com.google.testing.junit.runner.model.AntXmlResultWriter;
 import com.google.testing.junit.runner.model.XmlResultWriter;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.PrintStream;
 import java.text.DateFormat;
 import java.text.SimpleDateFormat;
@@ -67,11 +71,15 @@ public class BazelTestRunner {
    * <li>All tests pass: exit code of 0</li>
    * </ul>
    */
-  public static void main(String args[]) {
+  public static void main(String[] args) {
     PrintStream stderr = System.err;
 
     String suiteClassName = System.getProperty(TEST_SUITE_PROPERTY_NAME);
 
+    if (args.length >= 1 && args[args.length - 1].equals("--persistent_test_runner")) {
+      System.exit(runPersistentTestRunner(suiteClassName));
+    }
+
     if (!checkTestSuiteProperty(suiteClassName)) {
       System.exit(2);
     }
@@ -131,6 +139,7 @@ public class BazelTestRunner {
       }
     }
 
+    // TODO(kush): Use a new classloader for the following instantiation.
     JUnit4Runner runner =
         JUnit4Bazel.builder()
             .suiteClass(new SuiteClass(suite))
@@ -140,6 +149,48 @@ public class BazelTestRunner {
     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/rules/test/StandaloneTestStrategy.java b/src/main/java/com/google/devtools/build/lib/rules/test/StandaloneTestStrategy.java
index bbef175a7d..05eeb7abac 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/StandaloneTestStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/StandaloneTestStrategy.java
@@ -55,7 +55,7 @@ public class StandaloneTestStrategy extends TestStrategy {
   public static final String COLLECT_COVERAGE =
       "external/bazel_tools/tools/test/collect_coverage.sh";
 
-  private final Path workspace;
+  protected final Path workspace;
 
   public StandaloneTestStrategy(
       OptionsClassProvider requestOptions,
@@ -93,25 +93,6 @@ public class StandaloneTestStrategy extends TestStrategy {
     TestRunnerAction.ResolvedPaths resolvedPaths = action.resolve(execRoot);
     Map<String, String> env =
         getEnv(action, execRoot, runfilesDir, testTmpDir, resolvedPaths.getXmlOutputPath());
-
-    Map<String, String> info = new HashMap<>();
-
-    // This key is only understood by StandaloneSpawnStrategy.
-    info.put("timeout", "" + getTimeout(action));
-    info.putAll(action.getTestProperties().getExecutionInfo());
-
-    Artifact testSetup = action.getRuntimeArtifact(TEST_SETUP_BASENAME);
-    Spawn spawn =
-        new BaseSpawn(
-            getArgs(testSetup.getExecPathString(), COLLECT_COVERAGE, action),
-            env,
-            info,
-            new RunfilesSupplierImpl(
-                runfilesDir.asFragment(), action.getExecutionSettings().getRunfiles()),
-            action,
-            action.getTestProperties().getLocalResourceUsage(executionOptions.usingLocalTestJobs()),
-            ImmutableSet.of(resolvedPaths.getXmlOutputPath().relativeTo(execRoot)));
-
     Executor executor = actionExecutionContext.getExecutor();
 
     try {
@@ -128,7 +109,12 @@ public class StandaloneTestStrategy extends TestStrategy {
                       .getTestStderr());
           ResourceHandle handle = ResourceManager.instance().acquireResources(action, resources)) {
         TestResultData data =
-            execute(actionExecutionContext.withFileOutErr(fileOutErr), spawn, action);
+            execute(
+                actionExecutionContext.withFileOutErr(fileOutErr),
+                env,
+                action,
+                execRoot,
+                runfilesDir);
         appendStderr(fileOutErr.getOutputPath(), fileOutErr.getErrorPath());
         finalizeTest(actionExecutionContext, action, data);
       }
@@ -173,14 +159,35 @@ public class StandaloneTestStrategy extends TestStrategy {
     return vars;
   }
 
-  private TestResultData execute(
-      ActionExecutionContext actionExecutionContext, Spawn spawn, TestRunnerAction action)
-      throws ExecException, InterruptedException {
+  protected TestResultData execute(
+      ActionExecutionContext actionExecutionContext,
+      Map<String, String> environment,
+      TestRunnerAction action,
+      Path execRoot,
+      Path runfilesDir)
+      throws ExecException, InterruptedException, IOException {
     Executor executor = actionExecutionContext.getExecutor();
     Closeable streamed = null;
     Path testLogPath = action.getTestLog().getPath();
     TestResultData.Builder builder = TestResultData.newBuilder();
 
+    Map<String, String> info = new HashMap<>();
+    // This key is only understood by StandaloneSpawnStrategy.
+    info.put("timeout", "" + getTimeout(action));
+    info.putAll(action.getTestProperties().getExecutionInfo());
+
+    Artifact testSetup = action.getRuntimeArtifact(TEST_SETUP_BASENAME);
+    TestRunnerAction.ResolvedPaths resolvedPaths = action.resolve(execRoot);
+    Spawn spawn =
+        new BaseSpawn(
+            getArgs(testSetup.getExecPathString(), COLLECT_COVERAGE, action),
+            environment,
+            info,
+            new RunfilesSupplierImpl(
+                runfilesDir.asFragment(), action.getExecutionSettings().getRunfiles()),
+            action,
+            action.getTestProperties().getLocalResourceUsage(executionOptions.usingLocalTestJobs()),
+            ImmutableSet.of(resolvedPaths.getXmlOutputPath().relativeTo(execRoot)));
     long startTime = executor.getClock().currentTimeMillis();
     SpawnActionContext spawnActionContext = executor.getSpawnActionContext(action.getMnemonic());
     try {
diff --git a/src/main/java/com/google/devtools/build/lib/worker/BUILD b/src/main/java/com/google/devtools/build/lib/worker/BUILD
index 89b9776385..8e2c21fcb0 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/worker/BUILD
@@ -18,6 +18,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/sandbox",
         "//src/main/java/com/google/devtools/build/lib/standalone",
         "//src/main/java/com/google/devtools/common/options",
+        "//src/main/protobuf:test_status_java_proto",
         "//src/main/protobuf:worker_protocol_java_proto",
         "//third_party:apache_commons_pool2",
         "//third_party:guava",
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerActionContextProvider.java
index e8cd891060..e15af06b3e 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerActionContextProvider.java
@@ -14,10 +14,12 @@
 package com.google.devtools.build.lib.worker;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.devtools.build.lib.actions.ActionContextProvider;
 import com.google.devtools.build.lib.actions.Executor.ActionContext;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.exec.ExecutionOptions;
+import com.google.devtools.build.lib.rules.test.TestActionContext;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 
 /**
@@ -28,14 +30,21 @@ final class WorkerActionContextProvider extends ActionContextProvider {
 
   public WorkerActionContextProvider(
       CommandEnvironment env, BuildRequest buildRequest, WorkerPool workers) {
-    this.strategies =
-        ImmutableList.<ActionContext>of(
-            new WorkerSpawnStrategy(
-                env.getDirectories(),
-                workers,
-                buildRequest.getOptions(ExecutionOptions.class).verboseFailures,
-                buildRequest.getOptions(WorkerOptions.class).workerMaxRetries,
-                buildRequest.getOptions(WorkerOptions.class).workerVerbose));
+    int maxRetries = buildRequest.getOptions(WorkerOptions.class).workerMaxRetries;
+    ImmutableMultimap.Builder<String, String> extraFlags = ImmutableMultimap.builder();
+    extraFlags.putAll(buildRequest.getOptions(WorkerOptions.class).workerExtraFlags);
+
+    WorkerSpawnStrategy workerSpawnStrategy =
+        new WorkerSpawnStrategy(
+            env.getDirectories(),
+            workers,
+            buildRequest.getOptions(ExecutionOptions.class).verboseFailures,
+            maxRetries,
+            buildRequest.getOptions(WorkerOptions.class).workerVerbose,
+            extraFlags.build());
+    TestActionContext workerTestStrategy =
+        new WorkerTestStrategy(env, buildRequest, workers, maxRetries, extraFlags.build());
+    this.strategies = ImmutableList.of(workerSpawnStrategy, workerTestStrategy);
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerFilesHash.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerFilesHash.java
new file mode 100644
index 0000000000..08ab7f5e88
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerFilesHash.java
@@ -0,0 +1,41 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.worker;
+
+import com.google.common.hash.HashCode;
+import com.google.common.hash.Hasher;
+import com.google.common.hash.Hashing;
+import com.google.devtools.build.lib.actions.ActionExecutionContext;
+import com.google.devtools.build.lib.actions.ActionInput;
+import java.io.IOException;
+import java.nio.charset.Charset;
+
+/**
+ * Calculates the hash based on the files, which should be unchanged on disk for a worker to get
+ * reused.
+ */
+public class WorkerFilesHash {
+
+  public static HashCode getWorkerFilesHash(
+      Iterable<? extends ActionInput> toolFiles, ActionExecutionContext actionExecutionContext)
+      throws IOException {
+    Hasher hasher = Hashing.sha256().newHasher();
+    for (ActionInput tool : toolFiles) {
+      hasher.putString(tool.getExecPathString(), Charset.defaultCharset());
+      hasher.putBytes(actionExecutionContext.getActionInputFileCache().getDigest(tool));
+    }
+    return hasher.hash();
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
index d8fc26c776..ab3ffb88a1 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
@@ -13,9 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.lib.worker;
 
+import com.google.devtools.common.options.Converters;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.Options;
 import com.google.devtools.common.options.OptionsBase;
+import java.util.List;
+import java.util.Map.Entry;
 
 /**
  * Options related to worker processes.
@@ -73,6 +76,18 @@ public class WorkerOptions extends OptionsBase {
   public boolean workerVerbose;
 
   @Option(
+      name = "worker_extra_flag",
+      converter = Converters.AssignmentConverter.class,
+      defaultValue = "",
+      category = "strategy",
+      help =
+          "Extra command-flags that will be passed to worker processes in addition to "
+              + "--persistent_worker, keyed by mnemonic (e.g. --worker_extra_flag=Javac=--debug.",
+      allowMultiple = true
+  )
+  public List<Entry<String, String>> workerExtraFlags;
+
+  @Option(
     name = "worker_sandboxing",
     defaultValue = "false",
     category = "strategy",
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnStrategy.java
index 594258afe5..08d24b2652 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnStrategy.java
@@ -16,13 +16,13 @@ package com.google.devtools.build.lib.worker;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.base.Charsets;
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Multimap;
 import com.google.common.hash.HashCode;
-import com.google.common.hash.Hasher;
-import com.google.common.hash.Hashing;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputFileCache;
@@ -173,6 +173,7 @@ public final class WorkerSpawnStrategy implements SandboxedSpawnActionContext {
   private final Path execRoot;
   private final boolean verboseFailures;
   private final int maxRetries;
+  private final Multimap<String, String> extraFlags;
   private final boolean workerVerbose;
 
   public WorkerSpawnStrategy(
@@ -180,12 +181,15 @@ public final class WorkerSpawnStrategy implements SandboxedSpawnActionContext {
       WorkerPool workers,
       boolean verboseFailures,
       int maxRetries,
-      boolean workerVerbose) {
+      boolean workerVerbose,
+      Multimap<String, String> extraFlags) {
+    Preconditions.checkNotNull(workers);
     this.workers = Preconditions.checkNotNull(workers);
     this.execRoot = blazeDirs.getExecRoot();
     this.verboseFailures = verboseFailures;
     this.maxRetries = maxRetries;
     this.workerVerbose = workerVerbose;
+    this.extraFlags = extraFlags;
   }
 
   @Override
@@ -243,16 +247,21 @@ public final class WorkerSpawnStrategy implements SandboxedSpawnActionContext {
 
     FileOutErr outErr = actionExecutionContext.getFileOutErr();
 
-    ImmutableList<String> args = ImmutableList.<String>builder()
-        .addAll(spawn.getArguments().subList(0, spawn.getArguments().size() - 1))
-        .add("--persistent_worker")
-        .build();
+    ImmutableList<String> args =
+        ImmutableList.<String>builder()
+            .addAll(spawn.getArguments().subList(0, spawn.getArguments().size() - 1))
+            .add("--persistent_worker")
+            .addAll(
+                MoreObjects.firstNonNull(
+                    extraFlags.get(spawn.getMnemonic()), ImmutableList.<String>of()))
+            .build();
     ImmutableMap<String, String> env = spawn.getEnvironment();
 
     try {
       ActionInputFileCache inputFileCache = actionExecutionContext.getActionInputFileCache();
 
-      HashCode workerFilesHash = combineActionInputHashes(spawn.getToolFiles(), inputFileCache);
+      HashCode workerFilesHash = WorkerFilesHash.getWorkerFilesHash(
+          spawn.getToolFiles(), actionExecutionContext);
       Map<PathFragment, Path> inputFiles =
           new SpawnHelpers(execRoot).getMounts(spawn, actionExecutionContext);
       Set<PathFragment> outputFiles = SandboxHelpers.getOutputFiles(spawn);
@@ -329,17 +338,6 @@ public final class WorkerSpawnStrategy implements SandboxedSpawnActionContext {
     }
   }
 
-  private HashCode combineActionInputHashes(
-      Iterable<? extends ActionInput> toolFiles, ActionInputFileCache actionInputFileCache)
-      throws IOException {
-    Hasher hasher = Hashing.sha256().newHasher();
-    for (ActionInput tool : toolFiles) {
-      hasher.putString(tool.getExecPathString(), Charset.defaultCharset());
-      hasher.putBytes(actionInputFileCache.getDigest(tool));
-    }
-    return hasher.hash();
-  }
-
   private WorkResponse execInWorker(
       EventHandler eventHandler,
       WorkerKey key,
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
new file mode 100644
index 0000000000..c090c2e23c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerTestStrategy.java
@@ -0,0 +1,211 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.worker;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
+import com.google.common.hash.HashCode;
+import com.google.devtools.build.lib.actions.ActionExecutionContext;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.ExecException;
+import com.google.devtools.build.lib.actions.ExecutionStrategy;
+import com.google.devtools.build.lib.actions.Executor;
+import com.google.devtools.build.lib.actions.TestExecException;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.rules.test.StandaloneTestStrategy;
+import com.google.devtools.build.lib.rules.test.TestActionContext;
+import com.google.devtools.build.lib.rules.test.TestRunnerAction;
+import com.google.devtools.build.lib.runtime.CommandEnvironment;
+import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.view.test.TestStatus.BlazeTestStatus;
+import com.google.devtools.build.lib.view.test.TestStatus.TestCase;
+import com.google.devtools.build.lib.view.test.TestStatus.TestResultData;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import com.google.devtools.common.options.OptionsClassProvider;
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Runs TestRunnerAction actions in a worker. This is still experimental WIP.
+ * Do not use this strategy to run tests. <br>
+ *
+ * TODO(kush): List to things to cosider: <br>
+ * 1. Figure out if/how to honor the actions's execution info:
+ *      action.getTestProperties().getExecutionInfo() <br>
+ * 2. Figure out how to stream intermediate output when running in a Worker or block streamed
+ *      outputs for this strategy. <br>
+ * 3. Figure out how to add timeout facility. <br>
+ */
+@ExecutionStrategy(contextType = TestActionContext.class, name = { "experimental_worker" })
+public class WorkerTestStrategy extends StandaloneTestStrategy {
+  private final WorkerPool workerPool;
+  private final int maxRetries;
+  private final Multimap<String, String> extraFlags;
+
+  public WorkerTestStrategy(
+      CommandEnvironment env,
+      OptionsClassProvider requestOptions,
+      WorkerPool workerPool,
+      int maxRetries,
+      Multimap<String, String> extraFlags) {
+    super(
+        requestOptions,
+        env.getBlazeWorkspace().getBinTools(),
+        env.getClientEnv(),
+        env.getWorkspace());
+    this.workerPool = workerPool;
+    this.maxRetries = maxRetries;
+    this.extraFlags = extraFlags;
+  }
+
+  @Override
+  protected TestResultData execute(
+      ActionExecutionContext actionExecutionContext,
+      Map<String, String> environment,
+      TestRunnerAction action,
+      Path execRoot,
+      Path runfilesDir)
+      throws ExecException, InterruptedException, IOException {
+    List<String> startupArgs = getStartUpArgs(action);
+
+    return execInWorker(
+        actionExecutionContext,
+        environment,
+        action,
+        startupArgs,
+        execRoot,
+        maxRetries);
+  }
+
+  private TestResultData execInWorker(
+      ActionExecutionContext actionExecutionContext,
+      Map<String, String> environment,
+      TestRunnerAction action,
+      List<String> startupArgs,
+      Path execRoot,
+      int retriesLeft)
+      throws ExecException, InterruptedException, IOException {
+    Executor executor = actionExecutionContext.getExecutor();
+    TestResultData.Builder builder = TestResultData.newBuilder();
+
+    Path testLogPath = action.getTestLog().getPath();
+    Worker worker = null;
+    WorkerKey key = null;
+    long startTime = executor.getClock().currentTimeMillis();
+    try {
+      HashCode workerFilesHash = WorkerFilesHash.getWorkerFilesHash(
+          action.getTools(), actionExecutionContext);
+      key =
+          new WorkerKey(
+              startupArgs,
+              environment,
+              execRoot,
+              action.getMnemonic(),
+              workerFilesHash,
+              ImmutableMap.<PathFragment, Path>of(),
+              ImmutableSet.<PathFragment>of(),
+              /*mustBeSandboxed=*/false);
+      worker = workerPool.borrowObject(key);
+
+      WorkRequest request = WorkRequest.getDefaultInstance();
+      request.writeDelimitedTo(worker.getOutputStream());
+      worker.getOutputStream().flush();
+
+      WorkResponse response = WorkResponse.parseDelimitedFrom(worker.getInputStream());
+      actionExecutionContext.getFileOutErr().getErrorStream().write(
+          response.getOutputBytes().toByteArray());
+
+      long duration = executor.getClock().currentTimeMillis() - startTime;
+      builder.addTestTimes(duration);
+      builder.setRunDurationMillis(duration);
+      if (response.getExitCode() == 0) {
+        builder
+            .setTestPassed(true)
+            .setStatus(BlazeTestStatus.PASSED)
+            .setCachable(true)
+            .setPassedLog(testLogPath.getPathString());
+      } else {
+        builder
+            .setTestPassed(false)
+            .setStatus(BlazeTestStatus.FAILED)
+            .addFailedLogs(testLogPath.getPathString());
+      }
+      TestCase details = parseTestResult(
+          action.resolve(actionExecutionContext.getExecutor().getExecRoot()).getXmlOutputPath());
+      if (details != null) {
+        builder.setTestCase(details);
+      }
+
+      return builder.build();
+    } catch (IOException | InterruptedException e) {
+      if (e instanceof InterruptedException) {
+        // The user pressed Ctrl-C. Get out here quick.
+        retriesLeft = 0;
+      }
+
+      if (worker != null) {
+        workerPool.invalidateObject(key, worker);
+        worker = null;
+      }
+      if (retriesLeft > 0) {
+        // The worker process failed, but we still have some retries left. Let's retry with a fresh
+        // worker.
+        executor
+            .getEventHandler()
+            .handle(
+                Event.warn(
+                    key.getMnemonic()
+                        + " worker failed ("
+                        + e
+                        + "), invalidating and retrying with new worker..."));
+        return execInWorker(
+            actionExecutionContext,
+            environment,
+            action,
+            startupArgs,
+            execRoot,
+            retriesLeft - 1);
+      } else {
+        throw new TestExecException(e.getMessage());
+      }
+    } finally {
+      if (worker != null) {
+        workerPool.returnObject(key, worker);
+      }
+    }
+  }
+
+  private List<String> getStartUpArgs(TestRunnerAction action) throws ExecException {
+    Artifact testSetup = action.getRuntimeArtifact(TEST_SETUP_BASENAME);
+    List<String> args = getArgs(testSetup.getExecPathString(), "", action);
+    ImmutableList.Builder<String> startupArgs = ImmutableList.builder();
+    // Add test setup with no echo to prevent stdout corruption.
+    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()));
+    return startupArgs.build();
+  }
+}
diff --git a/src/test/shell/bazel/BUILD b/src/test/shell/bazel/BUILD
index e7cf472feb..b8ab2486d2 100644
--- a/src/test/shell/bazel/BUILD
+++ b/src/test/shell/bazel/BUILD
@@ -90,6 +90,13 @@ sh_test(
 )
 
 sh_test(
+    name = "persistent_test_runner_test",
+    size = "large",
+    srcs = ["persistent_test_runner_test.sh"],
+    data = [":test-deps"],
+)
+
+sh_test(
     name = "bazel_rules_test",
     size = "large",
     srcs = ["bazel_rules_test.sh"],
diff --git a/src/test/shell/bazel/persistent_test_runner_test.sh b/src/test/shell/bazel/persistent_test_runner_test.sh
new file mode 100755
index 0000000000..5a5d981d1f
--- /dev/null
+++ b/src/test/shell/bazel/persistent_test_runner_test.sh
@@ -0,0 +1,140 @@
+#!/bin/bash
+#
+# Copyright 2016 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Correctness tests for using a Persistent TestRunner.
+#
+
+# Load the test setup defined in the parent directory
+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 test_simple_scenario() {
+  mkdir -p java/testrunners || fail "mkdir failed"
+
+  cat > java/testrunners/TestsPass.java <<EOF
+package testrunners;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.junit.Test;
+
+@RunWith(JUnit4.class)
+public class TestsPass {
+
+  @Test
+  public void testPass() {
+    // This passes
+  }
+}
+EOF
+
+  cat > java/testrunners/TestsFail.java <<EOF
+package testrunners;
+import static org.junit.Assert.fail;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.junit.Test;
+
+@RunWith(JUnit4.class)
+public class TestsFail {
+
+  @Test
+  public void testFail() {
+    fail("Test is supposed to fail");
+  }
+}
+EOF
+
+  cat > java/testrunners/BUILD <<EOF
+java_test(name = "TestsPass",
+          srcs = ['TestsPass.java'],
+          deps = ['@bazel_tools//tools/jdk:TestRunner_deploy.jar'],
+)
+
+java_test(name = "TestsFail",
+          srcs = ['TestsFail.java'],
+          deps = ['@bazel_tools//tools/jdk:TestRunner_deploy.jar'],
+)
+EOF
+
+  bazel test --test_strategy=experimental_worker //java/testrunners:TestsPass \
+      || fail "Test fails unexpectedly"
+
+  bazel test --test_strategy=experimental_worker //java/testrunners:TestsFail \
+      && fail "Test passes unexpectedly" \
+      || true
+}
+
+# TODO(kush): Enable this test once we're able to reload modified classes in persistent test runner.
+function DISABLED_test_reload_modified_classes() {
+  mkdir -p java/testrunners || fail "mkdir failed"
+
+  # Create a passing test.
+  cat > java/testrunners/Tests.java <<EOF
+package testrunners;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.junit.Test;
+
+@RunWith(JUnit4.class)
+public class Tests {
+
+  @Test
+  public void testPass() {
+    // This passes
+  }
+}
+EOF
+
+  cat > java/testrunners/BUILD <<EOF
+java_test(name = "Tests",
+          srcs = ['Tests.java'],
+          deps = ['@bazel_tools//tools/jdk:TestRunner_deploy.jar'],
+)
+EOF
+
+  bazel test --test_strategy=experimental_worker //java/testrunners:Tests &> $TEST_log \
+      || fail "Test fails unexpectedly"
+
+  # Now get the test to fail.
+  cat > java/testrunners/Tests.java <<EOF
+package testrunners;
+import static org.junit.Assert.fail;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.junit.Test;
+
+@RunWith(JUnit4.class)
+public class Tests {
+
+  @Test
+  public void testPass() {
+    fail("Test is supposed to fail now");
+  }
+}
+EOF
+
+  bazel test --test_strategy=experimental_worker //java/testrunners:Tests &> $TEST_log \
+      && fail "Test passes unexpectedly" \
+      || true
+}
+
+
+run_suite "Persistent Test Runner tests"
\ No newline at end of file