Use zinc worker for DiscoverTestMain and GetTestTasks (#6125)

This should reduce the performance cost of spawning a new JVM in
scenarios where we already have a JVM worker running

Author: lihaoyi

ci/mill-bootstrap.patch

@@ -0,0 +1,14 @@
+diff --git a/integration/package.mill b/integration/package.mill
+index 5c388fafb6e..31763326edf 100644
+--- a/integration/package.mill
++++ b/integration/package.mill
+@@ -75,7 +75,8 @@ object `package` extends mill.Module {
+             testReportXml(),
+             javaHome().map(_.path),
+             testParallelism(),
+-            testLogLevel()
++            testLogLevel(),
++            jvmWorker = jvmWorker().worker()
+           )
+           testModuleUtil.runTests()
+         }

libs/javalib/api/src/mill/javalib/api/JvmWorkerApi.scala

@@ -47,4 +47,14 @@ trait JvmWorkerApi {
       javaHome: Option[os.Path],
       args: Seq[String]
   )(using ctx: JvmWorkerApi.Ctx): Boolean
+
+  def discoverTests(
+      value: mill.javalib.api.internal.ZincDiscoverTests,
+      javaHome: Option[os.Path]
+  )(using ctx: JvmWorkerApi.Ctx): Seq[String] = Nil
+
+  def getTestTasks(
+      value: mill.javalib.api.internal.ZincGetTestTasks,
+      javaHome: Option[os.Path]
+  )(using ctx: JvmWorkerApi.Ctx): Seq[String] = Nil
 }

libs/javalib/api/src/mill/javalib/api/internal/zinc_operations.scala

@@ -36,3 +36,14 @@ case class ZincScaladocJar(
     scalacPluginClasspath: Seq[PathRef],
     args: Seq[String]
 ) derives upickle.ReadWriter
+
+case class ZincDiscoverTests(runCp: Seq[os.Path], testCp: Seq[os.Path], framework: String)
+    derives upickle.ReadWriter
+
+case class ZincGetTestTasks(
+    runCp: Seq[os.Path],
+    testCp: Seq[os.Path],
+    framework: String,
+    selectors: Seq[String],
+    args: Seq[String]
+) derives upickle.ReadWriter

libs/javalib/src/mill/javalib/TestModule.scala

@@ -13,14 +13,8 @@ import mill.javalib.bsp.BspModule
 import mill.util.Jvm
 import mill.api.JsonFormatters.given
 import mill.constants.EnvVars
-import mill.javalib.testrunner.{
-  DiscoverTestsMain,
-  Framework,
-  TestArgs,
-  TestResult,
-  TestRunner,
-  TestRunnerUtils
-}
+import mill.javalib.api.internal.ZincDiscoverTests
+import mill.javalib.testrunner.{Framework, TestArgs, TestResult, TestRunner, TestRunnerUtils}
 
 import java.nio.file.Path
 
@@ -67,27 +61,15 @@ trait TestModule
    * Test classes (often called test suites) discovered by the configured [[testFramework]].
    */
   def discoveredTestClasses: T[Seq[String]] = Task {
-    val classes = if (javaHome().isDefined) {
-      Jvm.callProcess(
-        mainClass = "mill.javalib.testrunner.DiscoverTestsMain",
-        classPath = jvmWorker().scalalibClasspath().map(_.path).toVector,
-        mainArgs =
-          runClasspath().flatMap(p => Seq("--runCp", p.path.toString())) ++
-            testClasspath().flatMap(p => Seq("--testCp", p.path.toString())) ++
-            Seq("--framework", testFramework()),
-        javaHome = javaHome().map(_.path),
-        stdin = os.Inherit,
-        stdout = os.Pipe,
-        cwd = Task.dest
-      ).out.lines()
-    } else {
-      DiscoverTestsMain.main0(
+    val discoveredTests = jvmWorker().worker().discoverTests(
+      ZincDiscoverTests(
         runClasspath().map(_.path),
         testClasspath().map(_.path),
         testFramework()
-      )
-    }
-    classes.sorted
+      ),
+      javaHome().map(_.path)
+    )
+    discoveredTests.sorted
   }
 
   /**
@@ -266,7 +248,8 @@ trait TestModule
         javaHome().map(_.path),
         testParallelism(),
         testLogLevel(),
-        propagateEnv()
+        propagateEnv(),
+        jvmWorker().worker()
       )
       testModuleUtil.runTests()
     }

libs/javalib/src/mill/javalib/TestModuleUtil.scala

@@ -17,7 +17,8 @@ import mill.api.Logger
 
 import java.util.concurrent.ConcurrentHashMap
 import mill.api.BuildCtx
-import mill.javalib.testrunner.{GetTestTasksMain, TestArgs, TestResult, TestRunnerUtils}
+import mill.javalib.api.internal.ZincGetTestTasks
+import mill.javalib.testrunner.{TestArgs, TestResult, TestRunnerUtils}
 import os.Path
 
 import scala.annotation.unused
@@ -46,7 +47,8 @@ final class TestModuleUtil(
     javaHome: Option[os.Path],
     testParallelism: Boolean,
     testLogLevel: TestReporter.LogLevel,
-    propagateEnv: Boolean = true
+    propagateEnv: Boolean = true,
+    jvmWorker: mill.javalib.api.JvmWorkerApi
 )(using ctx: mill.api.TaskCtx) {
 
   private val (jvmArgs, props) = TestModuleUtil.loadArgsAndProps(useArgsFile, forkArgs)
@@ -78,32 +80,16 @@ final class TestModuleUtil(
         // test group requires spawning a JVM which can take 1+ seconds to realize there are no
         // tests to run and shut down
 
-        val discoveredTests = if (javaHome.isDefined) {
-          Jvm.callProcess(
-            mainClass = "mill.javalib.testrunner.GetTestTasksMain",
-            classPath = scalalibClasspath.map(_.path).toVector,
-            mainArgs =
-              (runClasspath ++ testrunnerEntrypointClasspath).flatMap(p =>
-                Seq("--runCp", p.path.toString)
-              ) ++
-                testClasspath.flatMap(p => Seq("--testCp", p.path.toString)) ++
-                Seq("--framework", testFramework) ++
-                selectors.flatMap(s => Seq("--selectors", s)) ++
-                args.flatMap(s => Seq("--args", s)),
-            javaHome = javaHome,
-            stdin = os.Inherit,
-            stdout = os.Pipe,
-            cwd = Task.dest
-          ).out.lines().toSet
-        } else {
-          GetTestTasksMain.main0(
+        val discoveredTests = jvmWorker.getTestTasks(
+          ZincGetTestTasks(
             (runClasspath ++ testrunnerEntrypointClasspath).map(_.path),
             testClasspath.map(_.path),
             testFramework,
             selectors,
             args
-          ).toSet
-        }
+          ),
+          javaHome
+        ).toSet
 
         filteredClassLists0.map(_.filter(discoveredTests)).filter(_.nonEmpty)
       }

libs/javalib/testrunner/src/mill/javalib/testrunner/DiscoverTestsMain.scala

@@ -3,13 +3,8 @@ package mill.javalib.testrunner
 import mill.api.daemon.internal.internal
 
 @internal object DiscoverTestsMain {
-  import mill.api.JsonFormatters.PathTokensReader
-
-  @mainargs.main
-  def main(runCp: Seq[os.Path], testCp: Seq[os.Path], framework: String): Unit = {
-    main0(runCp, testCp, framework).foreach(println)
-  }
-  def main0(runCp: Seq[os.Path], testCp: Seq[os.Path], framework: String): Seq[String] = {
+  def apply(args0: mill.javalib.api.internal.ZincDiscoverTests): Seq[String] = {
+    import args0.*
     mill.util.Jvm.withClassLoader(
       classPath = runCp,
       sharedPrefixes = Seq("sbt.testing.")
@@ -24,6 +19,4 @@ import mill.api.daemon.internal.internal
         }
     }
   }
-
-  def main(args: Array[String]): Unit = mainargs.Parser(this).runOrExit(args.toSeq)
 }

libs/javalib/testrunner/src/mill/javalib/testrunner/GetTestTasksMain.scala

@@ -3,25 +3,9 @@ package mill.javalib.testrunner
 import mill.api.daemon.internal.internal
 
 @internal object GetTestTasksMain {
-  import mill.api.JsonFormatters.PathTokensReader
-  @mainargs.main
-  def main(
-      runCp: Seq[os.Path],
-      testCp: Seq[os.Path],
-      framework: String,
-      selectors: Seq[String],
-      args: Seq[String]
-  ): Unit = {
-    main0(runCp, testCp, framework, selectors, args).foreach(println)
-  }
 
-  def main0(
-      runCp: Seq[os.Path],
-      testCp: Seq[os.Path],
-      framework: String,
-      selectors: Seq[String],
-      args: Seq[String]
-  ): Seq[String] = {
+  def apply(args0: mill.javalib.api.internal.ZincGetTestTasks): Seq[String] = {
+    import args0.*
     val globFilter = TestRunnerUtils.globFilter(selectors)
     mill.util.Jvm.withClassLoader(
       classPath = runCp,
@@ -38,6 +22,4 @@ import mill.api.daemon.internal.internal
         .toSeq
     }
   }
-
-  def main(args: Array[String]): Unit = mainargs.Parser(this).runOrExit(args.toSeq)
 }

libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala

@@ -85,6 +85,32 @@ class JvmWorkerImpl(args: JvmWorkerArgs) extends JvmWorkerApi with AutoCloseable
     result.result
   }
 
+  override def discoverTests(
+      op: mill.javalib.api.internal.ZincDiscoverTests,
+      javaHome: Option[os.Path]
+  )(using ctx: JvmWorkerApi.Ctx): Seq[String] = {
+    given RequestId = requestIds.next()
+
+    val zinc = zincApi(javaHome, JavaRuntimeOptions(Seq.empty))
+    val result = Timed(zinc.discoverTests(op))
+    fileLog(s"discoverTests took ${result.durationPretty}")
+    result.result
+  }
+
+  override def getTestTasks(
+      op: mill.javalib.api.internal.ZincGetTestTasks,
+      javaHome: Option[os.Path]
+  )(using
+      ctx: JvmWorkerApi.Ctx
+  ): Seq[String] = {
+    given RequestId = requestIds.next()
+
+    val zinc = zincApi(javaHome, JavaRuntimeOptions(Seq.empty))
+    val result = Timed(zinc.getTestTasks(op))
+    fileLog(s"getTestTasks took ${result.durationPretty}")
+    result.result
+  }
+
   override def close(): Unit = {
     zincLocalWorker.close()
     subprocessCache.close()
@@ -274,10 +300,7 @@ class JvmWorkerImpl(args: JvmWorkerArgs) extends JvmWorkerApi with AutoCloseable
   }
 
   /** Gives you API for the [[zincLocalWorker]] instance. */
-  private def localZincApi(
-      zincCtx: ZincWorker.InvocationContext,
-      log: Logger
-  ): ZincApi = {
+  private def localZincApi(zincCtx: ZincWorker.InvocationContext, log: Logger): ZincApi = {
     val zincDeps = ZincWorker.InvocationDependencies(
       log = log,
       consoleOut = ConsoleOut.printStreamOut(log.streams.err),
@@ -413,6 +436,20 @@ class JvmWorkerImpl(args: JvmWorkerArgs) extends JvmWorkerApi with AutoCloseable
           rpcClient(msg)
         }
       }
+
+      override def discoverTests(op: mill.javalib.api.internal.ZincDiscoverTests): Seq[String] = {
+        withRpcClient(serverRpcToClientHandler(reporter = None, log, cacheKey)) { rpcClient =>
+          val msg = ZincWorkerRpcServer.ClientToServer.DiscoverTests(op)
+          rpcClient(msg)
+        }
+      }
+
+      override def getTestTasks(op: mill.javalib.api.internal.ZincGetTestTasks): Seq[String] = {
+        withRpcClient(serverRpcToClientHandler(reporter = None, log, cacheKey)) { rpcClient =>
+          val msg = ZincWorkerRpcServer.ClientToServer.GetTestTasks(op)
+          rpcClient(msg)
+        }
+      }
     }
   }
 

libs/javalib/worker/src/mill/javalib/zinc/ZincApi.scala

@@ -26,4 +26,8 @@ trait ZincApi {
   def scaladocJar(
       op: ZincScaladocJar
   ): Boolean
+
+  def discoverTests(value: mill.javalib.api.internal.ZincDiscoverTests): Seq[String]
+
+  def getTestTasks(value: mill.javalib.api.internal.ZincGetTestTasks): Seq[String]
 }

libs/javalib/worker/src/mill/javalib/zinc/ZincWorker.scala

@@ -33,9 +33,7 @@ import java.util.Optional
 import scala.collection.mutable
 
 /** @param jobs number of parallel jobs */
-class ZincWorker(
-    jobs: Int
-) extends AutoCloseable { self =>
+class ZincWorker(jobs: Int) extends AutoCloseable { self =>
   private val incrementalCompiler = new sbt.internal.inc.IncrementalCompilerImpl()
   private val compilerBridgeLocks: mutable.Map[String, MemoryLock] = mutable.Map.empty
 
@@ -304,6 +302,12 @@ class ZincWorker(
 
     override def scaladocJar(op: ZincScaladocJar): Boolean =
       self.scaladocJar(op, deps.compilerBridge)
+
+    override def discoverTests(op: mill.javalib.api.internal.ZincDiscoverTests): Seq[String] =
+      mill.javalib.testrunner.DiscoverTestsMain(op)
+
+    override def getTestTasks(op: mill.javalib.api.internal.ZincGetTestTasks): Seq[String] =
+      mill.javalib.testrunner.GetTestTasksMain(op)
   }
 
   def close(): Unit = {