Add --experimental_persistent_test_runner for java tests.

RELNOTES: None.
PiperOrigin-RevId: 289618689

src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java

@@ -84,6 +84,7 @@ public final class RuleConfiguredTargetBuilder {
 
   private NestedSetBuilder<Artifact> filesToRunBuilder = NestedSetBuilder.stableOrder();
   private RunfilesSupport runfilesSupport;
+  private Runfiles persistentTestRunnerRunfiles;
   private Artifact executable;
   private ImmutableSet<ActionAnalysisMetadata> actionsWithoutExtraAction = ImmutableSet.of();
 
@@ -419,9 +420,9 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
     TestParams testParams =
         testActionBuilder
             .setFilesToRunProvider(filesToRunProvider)
+            .setPersistentTestRunnerRunfiles(persistentTestRunnerRunfiles)
             .setExecutionRequirements(
-                (ExecutionInfo) providersBuilder
-                    .getProvider(ExecutionInfo.PROVIDER.getKey()))
+                (ExecutionInfo) providersBuilder.getProvider(ExecutionInfo.PROVIDER.getKey()))
             .setShardCount(explicitShardCount)
             .build();
     ImmutableList<String> testTags = ImmutableList.copyOf(ruleContext.getRule().getRuleTags());
@@ -578,6 +579,11 @@ public RuleConfiguredTargetBuilder setRunfilesSupport(
     return this;
   }
 
+  public RuleConfiguredTargetBuilder setPersistentTestRunnerRunfiles(Runfiles testSupportRunfiles) {
+    this.persistentTestRunnerRunfiles = testSupportRunfiles;
+    return this;
+  }
+
   /**
    * Set the files to build.
    */

src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java

@@ -22,11 +22,13 @@
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
+import com.google.devtools.build.lib.actions.RunfilesSupplier;
 import com.google.devtools.build.lib.analysis.AnalysisEnvironment;
 import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.PrerequisiteArtifacts;
 import com.google.devtools.build.lib.analysis.RuleContext;
+import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesSupplierImpl;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.analysis.ShToolchain;
@@ -68,6 +70,7 @@ public final class TestActionBuilder {
 
   private final RuleContext ruleContext;
   private RunfilesSupport runfilesSupport;
+  private Runfiles persistentTestRunnerRunfiles;
   private Artifact executable;
   private ExecutionInfo executionRequirements;
   private InstrumentedFilesInfo instrumentedFiles;
@@ -120,6 +123,11 @@ public TestActionBuilder setFilesToRunProvider(FilesToRunProvider provider) {
     return this;
   }
 
+  public TestActionBuilder setPersistentTestRunnerRunfiles(Runfiles runfiles) {
+    this.persistentTestRunnerRunfiles = runfiles;
+    return this;
+  }
+
   public TestActionBuilder setInstrumentedFiles(@Nullable InstrumentedFilesInfo instrumentedFiles) {
     this.instrumentedFiles = instrumentedFiles;
     return this;
@@ -306,6 +314,7 @@ private TestParams createTestAction(int shards) {
               runfilesSupport,
               executable,
               instrumentedFileManifest,
+              /* persistentTestRunnerFlagFile= */ null,
               shards,
               runsPerTest);
       inputsBuilder.add(instrumentedFileManifest);
@@ -315,9 +324,16 @@ private TestParams createTestAction(int shards) {
         extraTestEnv.put(coverageEnvEntry.getFirst(), coverageEnvEntry.getSecond());
       }
     } else {
+      Artifact flagFile = null;
+      // The worker spawn runner expects a flag file containg the worker's flags.
+      if (testConfiguration.isPersistentTestRunner()) {
+        flagFile = ruleContext.getBinArtifact(ruleContext.getLabel().getName() + "_flag_file.txt");
+        inputsBuilder.add(flagFile);
+      }
+
       executionSettings =
           new TestTargetExecutionSettings(
-              ruleContext, runfilesSupport, executable, null, shards, runsPerTest);
+              ruleContext, runfilesSupport, executable, null, flagFile, shards, runsPerTest);
     }
 
     extraTestEnv.putAll(extraEnv);
@@ -369,11 +385,30 @@ private TestParams createTestAction(int shards) {
         boolean cancelConcurrentTests =
             testConfiguration.runsPerTestDetectsFlakes()
                 && testConfiguration.cancelConcurrentTests();
+        RunfilesSupplier testRunfilesSupplier;
+        if (testConfiguration.isPersistentTestRunner()) {
+          // Create a RunfilesSupplier from the persistent test runner's runfiles. Pass only the
+          // test runner's runfiles to avoid using a different worker for every test run.
+          testRunfilesSupplier =
+              new RunfilesSupplierImpl(
+                  /* runfilesDir= */ persistentTestRunnerRunfiles.getSuffix(),
+                  /* runfiles= */ persistentTestRunnerRunfiles,
+                  /* buildRunfileLinks= */ false,
+                  /* runfileLinksEnabled= */ false);
+        } else {
+          testRunfilesSupplier = RunfilesSupplierImpl.create(runfilesSupport);
+        }
+
+        ImmutableList.Builder<Artifact> tools = new ImmutableList.Builder<>();
+        if (testConfiguration.isPersistentTestRunner()) {
+          tools.add(testActionExecutable);
+          tools.add(executionSettings.getExecutable());
+        }
         TestRunnerAction testRunnerAction =
             new TestRunnerAction(
                 ruleContext.getActionOwner(),
                 inputs,
-                RunfilesSupplierImpl.create(runfilesSupport),
+                testRunfilesSupplier,
                 testActionExecutable,
                 isUsingTestWrapperInsteadOfTestSetupScript,
                 testXmlGeneratorExecutable,
@@ -392,7 +427,8 @@ private TestParams createTestAction(int shards) {
                         || executionSettings.needsShell(isExecutedOnWindows))
                     ? ShToolchain.getPathOrError(ruleContext)
                     : null,
-                cancelConcurrentTests);
+                cancelConcurrentTests,
+                tools.build());
 
         testOutputs.addAll(testRunnerAction.getSpawnOutputs());
         testOutputs.addAll(testRunnerAction.getOutputs());

src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java

@@ -155,6 +155,18 @@ public static class TestOptions extends FragmentOptions {
                 + "present. 'disabled' to never use test sharding.")
     public TestActionBuilder.TestShardingStrategy testShardingStrategy;
 
+    @Option(
+        name = "experimental_persistent_test_runner",
+        defaultValue = "false",
+        documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
+        effectTags = {OptionEffectTag.UNKNOWN},
+        help =
+            "Allows running java_test targets locally within a persistent worker. "
+                + "To enable the persistent test runner one must run bazel test with the flags:"
+                + "--test_strategy=local --strategy=TestRunner=worker "
+                + " --experimental_persistent_test_runner")
+    public boolean persistentTestRunner;
+
     @Option(
       name = "runs_per_test",
       allowMultiple = true,
@@ -313,6 +325,10 @@ public TestActionBuilder.TestShardingStrategy testShardingStrategy() {
     return options.testShardingStrategy;
   }
 
+  public boolean isPersistentTestRunner() {
+    return options.persistentTestRunner;
+  }
+
   public Label getCoverageSupport(){
     return options.coverageSupport;
   }

src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java

@@ -173,10 +173,11 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
       BuildConfiguration configuration,
       String workspaceName,
       @Nullable PathFragment shExecutable,
-      boolean cancelConcurrentTestsOnSuccess) {
+      boolean cancelConcurrentTestsOnSuccess,
+      Iterable<Artifact> tools) {
     super(
         owner,
-        /*tools=*/ NestedSetBuilder.emptySet(Order.STABLE_ORDER),
+        NestedSetBuilder.wrap(Order.STABLE_ORDER, tools),
         inputs,
         runfilesSupplier,
         nonNullAsSet(testLog, cacheStatus, coverageArtifact),
@@ -562,6 +563,11 @@ public void setupEnvVariables(Map<String, String> env, Duration timeout) {
       env.put("RUNFILES_MANIFEST_ONLY", "1");
     }
 
+    if (testConfiguration.isPersistentTestRunner()) {
+      // Let the test runner know it runs persistently within a worker.
+      env.put("PERSISTENT_TEST_RUNNER", "true");
+    }
+
     if (isCoverageMode()) {
       // Instruct remote-runtest.sh/local-runtest.sh not to cd into the runfiles directory.
       // TODO(ulfjack): Find a way to avoid setting this variable.
@@ -1138,3 +1144,4 @@ private ActionContinuationOrResult process(TestAttemptResult result, int actualM
     }
   }
 }
+

src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetExecutionSettings.java

@@ -14,15 +14,18 @@
 
 package com.google.devtools.build.lib.analysis.test;
 
+import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.CommandLine;
+import com.google.devtools.build.lib.actions.CommandLineExpansionException;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
+import com.google.devtools.build.lib.analysis.actions.FileWriteAction;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.RunUnder;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget.Mode;
@@ -54,6 +57,7 @@ public final class TestTargetExecutionSettings {
       RunfilesSupport runfilesSupport,
       Artifact executable,
       Artifact instrumentedFileManifest,
+      Artifact persistentTestRunnerFlagFile,
       int shards,
       int runs) {
     Preconditions.checkArgument(TargetUtils.isTestRule(ruleContext.getRule()));
@@ -62,8 +66,35 @@ public final class TestTargetExecutionSettings {
     TestConfiguration testConfig = config.getFragment(TestConfiguration.class);
 
     CommandLine targetArgs = runfilesSupport.getArgs();
-    testArguments =
-        CommandLine.concat(targetArgs, ImmutableList.copyOf(testConfig.getTestArguments()));
+    if (persistentTestRunnerFlagFile != null) {
+      // If an flag artifact exists for the persistent test runner, register an action that writes
+      // the test arguments to the said artifact. When the test runner runs in a persistent worker,
+      // the worker expects to find the test arguments in a special flag file.
+      ImmutableList.Builder<String> testTargetArgs = new ImmutableList.Builder<>();
+      try {
+        testTargetArgs.addAll(targetArgs.arguments());
+      } catch (CommandLineExpansionException e) {
+        // Don't fail the build and ignore the runfiles arguments.
+        ruleContext.ruleError("Could not expand test target arguments: " + e.getMessage());
+      }
+      testTargetArgs.addAll(testConfig.getTestArguments());
+      ruleContext.registerAction(
+          FileWriteAction.create(
+              ruleContext.getActionOwner(),
+              persistentTestRunnerFlagFile,
+              /* fileContents= */ Joiner.on(System.lineSeparator()).join(testTargetArgs.build()),
+              /* makeExecutable= */ false,
+              /* allowCompression= */ FileWriteAction.Compression.DISALLOW));
+
+      // When using the persistent test runner the test arguments are passed through --flagfile.
+      testArguments =
+          CommandLine.of(
+              ImmutableList.of(
+                  "--flagfile=" + persistentTestRunnerFlagFile.getRootRelativePathString()));
+    } else {
+      testArguments =
+          CommandLine.concat(targetArgs, ImmutableList.copyOf(testConfig.getTestArguments()));
+    }
 
     totalShards = shards;
     totalRuns = runs;

src/main/java/com/google/devtools/build/lib/bazel/rules/java/BazelJavaSemantics.java

@@ -306,7 +306,42 @@ public Artifact createStubAction(
     }
     NestedSet<Artifact> classpath = classpathBuilder.build();
 
-    arguments.add(new ComputedClasspathSubstitution(classpath, workspacePrefix, isRunfilesEnabled));
+    if (JavaSemantics.isPersistentTestRunner(ruleContext)) {
+      // Create an artifact that stores the test's runtime classpath (excluding the test support
+      // classpath). The file is read by the test runner. The jars inside the file are loaded
+      // dynamically for every test run into a custom classloader.
+      arguments.add(
+          new ComputedClasspathSubstitution(
+              JavaSemantics.getTestSupportRuntimeClasspath(ruleContext),
+              workspacePrefix,
+              isRunfilesEnabled));
+
+      // Create an artifact that stores the test's runtime classpath.
+      Artifact testRuntimeClasspathArtifact =
+          ruleContext.getBinArtifact(
+              ruleContext.getLabel().getName() + "_test_runtime_classpath.txt");
+      ruleContext.registerAction(
+          new LazyWritePathsFileAction(
+              ruleContext.getActionOwner(),
+              testRuntimeClasspathArtifact,
+              javaCommon.getRuntimeClasspath(),
+              true));
+      filesBuilder.add(testRuntimeClasspathArtifact);
+
+      arguments.add(
+          Substitution.of(
+              TEST_RUNTIME_CLASSPATH_FILE_PLACEHOLDER,
+              "export WORKSPACE_PREFIX="
+                  + workspacePrefix
+                  + "\nexport TEST_RUNTIME_CLASSPATH_FILE=${JAVA_RUNFILES}"
+                  + File.separator
+                  + workspacePrefix
+                  + testRuntimeClasspathArtifact.getRootRelativePathString()));
+    } else {
+      arguments.add(
+          new ComputedClasspathSubstitution(classpath, workspacePrefix, isRunfilesEnabled));
+      arguments.add(Substitution.of(TEST_RUNTIME_CLASSPATH_FILE_PLACEHOLDER, ""));
+    }
 
     if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
       if (createCoverageMetadataJar) {
@@ -453,9 +488,14 @@ public void collectTargetsTreatedAsDeps(
       // targets may break, we are keeping it behind this flag.
       return;
     }
-    TransitiveInfoCollection testSupport = JavaSemantics.getTestSupport(ruleContext);
-    if (testSupport != null) {
-      builder.add(testSupport);
+    if (!JavaSemantics.isPersistentTestRunner(ruleContext)) {
+      // Only add the test support to the dependencies when running in regular mode.
+      // In persistent test runner mode don't pollute the classpath of the test with
+      // the test support classes.
+      TransitiveInfoCollection testSupport = JavaSemantics.getTestSupport(ruleContext);
+      if (testSupport != null) {
+        builder.add(testSupport);
+      }
     }
   }
 
@@ -717,3 +757,4 @@ public Artifact getObfuscatedConstantStringMap(RuleContext ruleContext)
   @Override
   public void checkDependencyRuleKinds(RuleContext ruleContext) {}
 }
+

src/main/java/com/google/devtools/build/lib/bazel/rules/java/java_stub_template.txt

@@ -106,6 +106,11 @@ for ARG in "$@"; do
   fi
 done
 
+# A file storing the test's runtime classpath, excluding the test support.
+# The file is read by the persistent test runner. The jars inside the file are
+# loaded dynamically for every test run into a custom classloader.
+%test_runtime_classpath_file%
+
 # Find our runfiles tree.  We need this to construct the classpath
 # (unless --singlejar was passed).
 #

src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java

@@ -37,6 +37,7 @@
 import com.google.devtools.build.lib.actions.TestExecException;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.test.TestActionContext;
+import com.google.devtools.build.lib.analysis.test.TestConfiguration;
 import com.google.devtools.build.lib.analysis.test.TestResult;
 import com.google.devtools.build.lib.analysis.test.TestRunnerAction;
 import com.google.devtools.build.lib.analysis.test.TestRunnerAction.ResolvedPaths;
@@ -119,6 +120,9 @@ public TestRunnerSpawn createTestRunnerSpawn(
       executionInfo.put(ExecutionRequirements.NO_CACHE, "");
     }
     executionInfo.put(ExecutionRequirements.TIMEOUT, "" + getTimeout(action).getSeconds());
+    if (action.getConfiguration().getFragment(TestConfiguration.class).isPersistentTestRunner()) {
+      executionInfo.put(ExecutionRequirements.SUPPORTS_WORKERS, "1");
+    }
 
     ResourceSet localResourceUsage =
         action
@@ -135,7 +139,9 @@ public TestRunnerSpawn createTestRunnerSpawn(
             action.getRunfilesSupplier(),
             ImmutableMap.of(),
             /*inputs=*/ action.getInputs(),
-            /*tools=*/ NestedSetBuilder.emptySet(Order.STABLE_ORDER),
+            action.getConfiguration().getFragment(TestConfiguration.class).isPersistentTestRunner()
+                ? NestedSetBuilder.wrap(Order.STABLE_ORDER, action.getTools())
+                : NestedSetBuilder.emptySet(Order.STABLE_ORDER),
             ImmutableSet.copyOf(action.getSpawnOutputs()),
             localResourceUsage);
     return new StandaloneTestRunnerSpawn(

src/main/java/com/google/devtools/build/lib/rules/java/JavaBinary.java

@@ -337,6 +337,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
     Runfiles defaultRunfiles = runfilesBuilder.build();
 
     RunfilesSupport runfilesSupport = null;
+    Runfiles persistentTestRunnerRunfiles = null;
     NestedSetBuilder<Artifact> extraFilesToRunBuilder = NestedSetBuilder.stableOrder();
     if (createExecutable) {
       List<String> extraArgs =
@@ -350,6 +351,9 @@ public ConfiguredTarget create(RuleContext ruleContext)
           RunfilesSupport.withExecutable(
               ruleContext, defaultRunfiles, executableForRunfiles, extraArgs);
       extraFilesToRunBuilder.add(runfilesSupport.getRunfilesMiddleman());
+      if (JavaSemantics.isPersistentTestRunner(ruleContext)) {
+        persistentTestRunnerRunfiles = JavaSemantics.getTestSupportRunfiles(ruleContext);
+      }
     }
 
     RunfilesProvider runfilesProvider =
@@ -473,6 +477,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
         // shell script), on Windows they are different (the executable to run is a batch file, the
         // executable for runfiles is the shell script).
         .setRunfilesSupport(runfilesSupport, executableToRun)
+        .setPersistentTestRunnerRunfiles(persistentTestRunnerRunfiles)
         .addFilesToRun(extraFilesToRunBuilder.build())
         .add(
             JavaRuntimeClasspathProvider.class,