Fix docjar context classloader conflict in scala 3 (#4331)

Fixes #4316 This isn't a general solution to make context classloaders behave properly in the presence of long-lived worker objects instantiated from classloaders. Rather, we just hard-code support in the Scala3 docjar callsite to properly set the context classloader there, since these context-classloader-sensitive pieces of code are pretty rare in Mill and so we can just deal with them individually as they pop up CC @joan38
com-lihaoyi · Jan 15, 2025 · 2a40b17 · 2a40b17
1 parent c3faa7c
commit 2a40b17
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 15 deletions.
diff --git a/example/extending/jvmcode/3-worker/build.mill b/example/extending/jvmcode/3-worker/build.mill
@@ -27,9 +27,7 @@ trait GroovyGenerateJavaModule extends JavaModule {
   def groovyScript = Task.Source(millSourcePath / "generate.groovy")
 
   def groovyGeneratedResources = Task {
-    val oldCl = Thread.currentThread().getContextClassLoader
-    Thread.currentThread().setContextClassLoader(groovyWorker())
-    try {
+    mill.api.ClassLoader.withContextClassLoader(groovyWorker()) {
       groovyWorker()
         .loadClass("groovy.ui.GroovyMain")
         .getMethod("main", classOf[Array[String]])
@@ -41,7 +39,7 @@ trait GroovyGenerateJavaModule extends JavaModule {
             (Task.dest / "groovy-generated.html").toString
           )
         )
-    } finally Thread.currentThread().setContextClassLoader(oldCl)
+    }
     PathRef(Task.dest)
   }
 

diff --git a/main/api/src/mill/api/ClassLoader.scala b/main/api/src/mill/api/ClassLoader.scala
@@ -8,6 +8,14 @@ import java.net.{URL, URLClassLoader}
  */
 object ClassLoader {
 
+  def withContextClassLoader[T](cl: java.lang.ClassLoader)(t: => T): T = {
+    val thread = Thread.currentThread()
+    val oldCl = thread.getContextClassLoader()
+    try {
+      thread.setContextClassLoader(cl)
+      t
+    } finally thread.setContextClassLoader(oldCl)
+  }
   def java9OrAbove: Boolean = !System.getProperty("java.specification.version").startsWith("1.")
 
   def create(

diff --git a/runner/src/mill/runner/MillBuildBootstrap.scala b/runner/src/mill/runner/MillBuildBootstrap.scala
@@ -399,17 +399,15 @@ object MillBuildBootstrap {
       selectiveExecution: Boolean
   ): (Either[String, Seq[Any]], Seq[Watchable], Seq[Watchable]) = {
     rootModule.evalWatchedValues.clear()
-    val previousClassloader = Thread.currentThread().getContextClassLoader
     val evalTaskResult =
-      try {
-        Thread.currentThread().setContextClassLoader(rootModule.getClass.getClassLoader)
+      mill.api.ClassLoader.withContextClassLoader(rootModule.getClass.getClassLoader) {
         RunScript.evaluateTasksNamed(
           evaluator,
           targetsAndParams,
           SelectMode.Separated,
           selectiveExecution = selectiveExecution
         )
-      } finally Thread.currentThread().setContextClassLoader(previousClassloader)
+      }
 
     val moduleWatched = rootModule.watchedValues.toVector
     val addedEvalWatched = rootModule.evalWatchedValues.toVector

diff --git a/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala b/scalalib/worker/src/mill/scalalib/worker/ZincWorkerImpl.scala
@@ -235,13 +235,17 @@ class ZincWorkerImpl(
         val hasErrorsMethod = reporter.getClass().getMethod("hasErrors")
         !hasErrorsMethod.invoke(reporter).asInstanceOf[Boolean]
       } else if (ZincWorkerUtil.isScala3(scalaVersion)) {
-        val scaladocClass =
-          compilers.scalac().scalaInstance().loader().loadClass("dotty.tools.scaladoc.Main")
-        val scaladocMethod = scaladocClass.getMethod("run", classOf[Array[String]])
-        val reporter =
-          scaladocMethod.invoke(scaladocClass.getConstructor().newInstance(), args.toArray)
-        val hasErrorsMethod = reporter.getClass().getMethod("hasErrors")
-        !hasErrorsMethod.invoke(reporter).asInstanceOf[Boolean]
+        // DottyDoc makes use of `com.fasterxml.jackson.databind.Module` which
+        // requires the ContextClassLoader to be set appropriately
+        mill.api.ClassLoader.withContextClassLoader(getClass.getClassLoader) {
+          val scaladocClass =
+            compilers.scalac().scalaInstance().loader().loadClass("dotty.tools.scaladoc.Main")
+          val scaladocMethod = scaladocClass.getMethod("run", classOf[Array[String]])
+          val reporter =
+            scaladocMethod.invoke(scaladocClass.getConstructor().newInstance(), args.toArray)
+          val hasErrorsMethod = reporter.getClass().getMethod("hasErrors")
+          !hasErrorsMethod.invoke(reporter).asInstanceOf[Boolean]
+        }
       } else {
         val scaladocClass =
           compilers.scalac().scalaInstance().loader().loadClass("scala.tools.nsc.ScalaDoc")