Generalize Zinc worker

- Compiler bridges can now be either pre-compiled or on-demand-compiled
- Scala library/compiler jar discovery is now configurable
- Zinc compiler cache is now configurable, rather than being hardcoded at n=1

1 files changed, 133 insertions, 97 deletions

diff --git a/scalalib/worker/src/ZincWorkerImpl.scala b/scalalib/worker/src/ZincWorkerImpl.scala
index 705d4682..0c1201a9 100644
--- a/scalalib/worker/src/ZincWorkerImpl.scala
+++ b/scalalib/worker/src/ZincWorkerImpl.scala
@@ -1,5 +1,7 @@
 package mill.scalalib.worker
 
+import mill.scalalib.worker._
+
 import java.io.File
 import java.util.Optional
 
@@ -19,10 +21,53 @@ case class MockedLookup(am: File => Optional[CompileAnalysis]) extends PerClassp
     Locate.definesClass(classpathEntry)
 }
 
-class ZincWorkerImpl(ctx0: mill.api.Ctx,
-                     compilerBridgeClasspath: Array[String]) extends mill.scalalib.api.ZincWorkerApi{
+trait KeyedCache[T]{
+  def withCachedValue[V](key: Long)(f: => T)(f2: T => V): V
+}
+
+object KeyedCache{
+  class RandomBoundedCache[T](hotParallelism: Int, coldCacheSize: Int) extends KeyedCache[T]{
+    private[this] val random = new scala.util.Random(313373)
+    val available = new java.util.concurrent.Semaphore(hotParallelism)
+
+    // Awful asymptotic complexity, but our caches are tiny n < 10 so it doesn't matter
+    var cache = Array.fill[Option[(Long, T)]](coldCacheSize)(None)
+
+    def withCachedValue[V](key: Long)(f: => T)(f2: T => V): V = {
+      available.acquire()
+      val pickedValue = synchronized{
+        cache.indexWhere(_.exists(_._1 == key)) match {
+          case -1 => f
+          case i =>
+            val (k, v) = cache(i).get
+            cache(i) = None
+            v
+        }
+      }
+      val result = f2(pickedValue)
+      synchronized{
+        cache.indexWhere(_.isEmpty) match{
+          // Random eviction #YOLO
+          case -1 => cache(random.nextInt(cache.length)) = Some((key, pickedValue))
+          case i => cache(i) = Some((key, pickedValue))
+        }
+      }
+
+      available.release()
+      result
+    }
+  }
+}
+
+object ZincWorkerImpl{
+  type Ctx = mill.api.Ctx.Dest with mill.api.Ctx.Log with mill.api.Ctx.Home
+}
+class ZincWorkerImpl(compilerBridge: Either[(ZincWorkerImpl.Ctx, Array[os.Path], String => os.Path), String => os.Path],
+                     libraryJarNameGrep: (Agg[os.Path], String) => os.Path = grepJar(_, "scala-library", _),
+                     compilerJarNameGrep: (Agg[os.Path], String) => os.Path = grepJar(_, "scala-compiler", _),
+                     compilerCache: KeyedCache[Compilers] = new KeyedCache.RandomBoundedCache(1, 1)) {
   private val ic = new sbt.internal.inc.IncrementalCompilerImpl()
-  val javaOnlyCompilers = {
+  lazy val javaOnlyCompilers = {
     // Keep the classpath as written by the user
     val classpathOptions = ClasspathOptions.of(false, false, false, false, false)
 
@@ -42,61 +87,60 @@ class ZincWorkerImpl(ctx0: mill.api.Ctx,
     )
   }
 
-  @volatile var mixedCompilersCache = Option.empty[(Long, Compilers)]
-
   def docJar(scalaVersion: String,
-             compilerBridgeSources: os.Path,
              compilerClasspath: Agg[os.Path],
              scalacPluginClasspath: Agg[os.Path],
              args: Seq[String])
-            (implicit ctx: mill.api.Ctx): Boolean = {
-    val compilers: Compilers = prepareCompilers(
+            (implicit ctx: ZincWorkerImpl.Ctx): Boolean = {
+    withCompilers(
       scalaVersion,
-      compilerBridgeSources,
       compilerClasspath,
-      scalacPluginClasspath
-    )
-    val scaladocClass = compilers.scalac().scalaInstance().loader().loadClass("scala.tools.nsc.ScalaDoc")
-    val scaladocMethod = scaladocClass.getMethod("process", classOf[Array[String]])
-    scaladocMethod.invoke(scaladocClass.newInstance(), args.toArray).asInstanceOf[Boolean]
+      scalacPluginClasspath,
+    ) { compilers: Compilers =>
+      val scaladocClass = compilers.scalac().scalaInstance().loader().loadClass("scala.tools.nsc.ScalaDoc")
+      val scaladocMethod = scaladocClass.getMethod("process", classOf[Array[String]])
+      scaladocMethod.invoke(scaladocClass.newInstance(), args.toArray).asInstanceOf[Boolean]
+    }
   }
   /** Compile the bridge if it doesn't exist yet and return the output directory.
-   *  TODO: Proper invalidation, see #389
-   */
-  def compileZincBridgeIfNeeded(scalaVersion: String,
-                                sourcesJar: os.Path,
-                                compilerJars: Array[File]): os.Path = {
-    val workingDir = ctx0.dest / scalaVersion
-    val compiledDest = workingDir / 'compiled
-    if (!os.exists(workingDir)) {
-
-      ctx0.log.info("Compiling compiler interface...")
-
-      os.makeDir.all(workingDir)
-      os.makeDir.all(compiledDest)
-
-      val sourceFolder = mill.api.IO.unpackZip(sourcesJar)(workingDir)
-      val classloader = mill.api.ClassLoader.create(compilerJars.map(_.toURI.toURL), null)(ctx0)
-      val compilerMain = classloader.loadClass(
-        if (isDotty(scalaVersion))
-          "dotty.tools.dotc.Main"
-        else
-          "scala.tools.nsc.Main"
-      )
-      val argsArray = Array[String](
-        "-d", compiledDest.toString,
-        "-classpath", (compilerJars ++ compilerBridgeClasspath).mkString(File.pathSeparator)
-      ) ++ os.walk(sourceFolder.path).filter(_.ext == "scala").map(_.toString)
-
-      compilerMain.getMethod("process", classOf[Array[String]])
-        .invoke(null, argsArray)
+    *  TODO: Proper invalidation, see #389
+    */
+  def compileZincBridgeIfNeeded(scalaVersion: String, compilerJars: Array[File]): os.Path = {
+    compilerBridge match{
+      case Right(compiled) => compiled(scalaVersion)
+      case Left((ctx0, compilerBridgeClasspath, srcJars)) =>
+        val workingDir = ctx0.dest / scalaVersion
+        val compiledDest = workingDir / 'compiled
+        if (!os.exists(workingDir)) {
+          ctx0.log.info("Compiling compiler interface...")
+
+          os.makeDir.all(workingDir)
+          os.makeDir.all(compiledDest)
+
+          val sourceFolder = mill.api.IO.unpackZip(srcJars(scalaVersion))(workingDir)
+          val classloader = mill.api.ClassLoader.create(compilerJars.map(_.toURI.toURL), null)(ctx0)
+          val compilerMain = classloader.loadClass(
+            if (isDotty(scalaVersion)) "dotty.tools.dotc.Main"
+            else "scala.tools.nsc.Main"
+          )
+          val argsArray = Array[String](
+            "-d", compiledDest.toString,
+            "-Ylog-classpath",
+            "-classpath", (compilerJars ++ compilerBridgeClasspath).mkString(File.pathSeparator)
+          ) ++ os.walk(sourceFolder.path).filter(_.ext == "scala").map(_.toString)
+
+          compilerMain.getMethod("process", classOf[Array[String]])
+            .invoke(null, argsArray)
+        }
+        compiledDest
     }
-    compiledDest
+
   }
 
 
 
-  def discoverMainClasses(compilationResult: CompilationResult)(implicit ctx: mill.api.Ctx): Seq[String] = {
+  def discoverMainClasses(compilationResult: CompilationResult)
+                         (implicit ctx: ZincWorkerImpl.Ctx): Seq[String] = {
     def toScala[A](o: Optional[A]): Option[A] = if (o.isPresent) Some(o.get) else None
 
     toScala(FileAnalysisStore.binary(compilationResult.analysisFile.toIO).get())
@@ -114,7 +158,7 @@ class ZincWorkerImpl(ctx0: mill.api.Ctx,
                   sources: Agg[os.Path],
                   compileClasspath: Agg[os.Path],
                   javacOptions: Seq[String])
-                 (implicit ctx: mill.api.Ctx): mill.api.Result[CompilationResult] = {
+                 (implicit ctx: ZincWorkerImpl.Ctx): mill.api.Result[CompilationResult] = {
     compileInternal(
       upstreamCompileOutput,
       sources,
@@ -131,72 +175,63 @@ class ZincWorkerImpl(ctx0: mill.api.Ctx,
                    javacOptions: Seq[String],
                    scalaVersion: String,
                    scalacOptions: Seq[String],
-                   compilerBridgeSources: os.Path,
                    compilerClasspath: Agg[os.Path],
                    scalacPluginClasspath: Agg[os.Path])
-                  (implicit ctx: mill.api.Ctx): mill.api.Result[CompilationResult] = {
-    val compilers: Compilers = prepareCompilers(
+                  (implicit ctx: ZincWorkerImpl.Ctx): mill.api.Result[CompilationResult] = {
+    withCompilers(
       scalaVersion,
-      compilerBridgeSources,
       compilerClasspath,
-      scalacPluginClasspath
-    )
-
-    compileInternal(
-      upstreamCompileOutput,
-      sources,
-      compileClasspath,
-      javacOptions,
-      scalacOptions = scalacPluginClasspath.map(jar => s"-Xplugin:${jar}").toSeq ++ scalacOptions,
-      compilers
-    )
+      scalacPluginClasspath,
+    ) {compilers: Compilers =>
+      compileInternal(
+        upstreamCompileOutput,
+        sources,
+        compileClasspath,
+        javacOptions,
+        scalacOptions = scalacPluginClasspath.map(jar => s"-Xplugin:$jar").toSeq ++ scalacOptions,
+        compilers
+      )
+    }
   }
 
-  private def prepareCompilers(scalaVersion: String,
-                               compilerBridgeSources: os.Path,
+  private def withCompilers[T](scalaVersion: String,
                                compilerClasspath: Agg[os.Path],
                                scalacPluginClasspath: Agg[os.Path])
-                              (implicit ctx: mill.api.Ctx)= {
+                              (f: Compilers => T)
+                              (implicit ctx: ZincWorkerImpl.Ctx)= {
     val combinedCompilerClasspath = compilerClasspath ++ scalacPluginClasspath
     val combinedCompilerJars = combinedCompilerClasspath.toArray.map(_.toIO)
 
-    val compilerBridge = compileZincBridgeIfNeeded(
-      scalaVersion,
-      compilerBridgeSources,
-      compilerClasspath.toArray.map(_.toIO)
-    )
-    val compilerBridgeSig = os.mtime(compilerBridge)
+    val compiledCompilerBridge =
+      compileZincBridgeIfNeeded(scalaVersion, compilerClasspath.toArray.map(_.toIO))
+
+    val compilerBridgeSig = os.mtime(compiledCompilerBridge)
 
     val compilersSig =
       compilerBridgeSig +
         combinedCompilerClasspath.map(p => p.toString().hashCode + os.mtime(p)).sum
 
-    val compilers = mixedCompilersCache match {
-      case Some((k, v)) if k == compilersSig => v
-      case _ =>
-        val compilerName =
-          if (isDotty(scalaVersion))
-            s"dotty-compiler_${scalaBinaryVersion(scalaVersion)}"
-          else
-            "scala-compiler"
-        val scalaInstance = new ScalaInstance(
-          version = scalaVersion,
-          loader = mill.api.ClassLoader.create(combinedCompilerJars.map(_.toURI.toURL), null),
-          libraryJar = grepJar(compilerClasspath, "scala-library", scalaVersion).toIO,
-          compilerJar = grepJar(compilerClasspath, compilerName, scalaVersion).toIO,
-          allJars = combinedCompilerJars,
-          explicitActual = None
-        )
-        val compilers = ic.compilers(
-          scalaInstance,
-          ClasspathOptionsUtil.boot,
-          None,
-          ZincUtil.scalaCompiler(scalaInstance, compilerBridge.toIO)
-        )
-        mixedCompilersCache = Some((compilersSig, compilers))
-        compilers
-    }
-    compilers
+    compilerCache.withCachedValue(compilersSig){
+      val compilerJar =
+        if (isDotty(scalaVersion))
+          grepJar(compilerClasspath, s"dotty-compiler_${scalaBinaryVersion(scalaVersion)}", scalaVersion)
+        else
+          compilerJarNameGrep(compilerClasspath, scalaVersion)
+      val scalaInstance = new ScalaInstance(
+        version = scalaVersion,
+        loader = mill.api.ClassLoader.create(combinedCompilerJars.map(_.toURI.toURL), null),
+        libraryJar = libraryJarNameGrep(compilerClasspath, scalaVersion).toIO,
+        compilerJar = compilerJar.toIO,
+        allJars = combinedCompilerJars,
+        explicitActual = None
+      )
+      ic.compilers(
+        scalaInstance,
+        ClasspathOptionsUtil.boot,
+        None,
+        ZincUtil.scalaCompiler(scalaInstance, compiledCompilerBridge.toIO)
+      )
+    }(f)
   }
 
   private def compileInternal(upstreamCompileOutput: Seq[CompilationResult],
@@ -205,9 +240,10 @@ class ZincWorkerImpl(ctx0: mill.api.Ctx,
                               javacOptions: Seq[String],
                               scalacOptions: Seq[String],
                               compilers: Compilers)
-                             (implicit ctx: mill.api.Ctx): mill.api.Result[CompilationResult] = {
+                             (implicit ctx: ZincWorkerImpl.Ctx): mill.api.Result[CompilationResult] = {
     os.makeDir.all(ctx.dest)
 
+    println("compileClasspath " + compileClasspath)
     val logger = {
       val consoleAppender = MainAppender.defaultScreen(ConsoleOut.printStreamOut(
         ctx.log.outputStream