Close ExternalSummary class loaders right after use (#5515)

It seems the class loaders used by ExternalSummary are only used right
after having been created, and aren't needed anymore after that. So the
changes here make sure we do close them right after.

Author: alexarchambault

integration/feature/leak-hygiene/src/LeakHygieneTests.scala

@@ -14,12 +14,14 @@ import utest._
  */
 object LeakHygieneTests extends UtestIntegrationTestSuite {
   def checkClassloaders(tester: IntegrationTester)(kvs: (String, Int)*) = {
-    val res = tester.eval(("show", "countClassLoaders"))
+    val res = tester.eval(("show", "countClassLoaders"), check = true)
 
     val read = upickle.default.read[SortedMap[String, Int]](res.out)
     val expected = SortedMap(kvs*)
-    // pprint.log(read)
-    // pprint.log(expected)
+    if (read != expected) {
+      pprint.log(expected)
+      pprint.log(read)
+    }
     assert(read == expected)
   }
 
@@ -64,7 +66,6 @@ object LeakHygieneTests extends UtestIntegrationTestSuite {
         mill.constants.DebugLog("\nstart")
         checkClassloaders(tester)(
           "mill.daemon.MillBuildBootstrap#processRunClasspath classLoader cl" -> 1,
-          "mill.codesig.ExternalSummary.apply upstreamClassloader" -> 1,
           "mill.javalib.JvmWorkerModule#worker cl" -> 1,
           "mill.javalib.worker.JvmWorkerImpl#scalaCompilerCache.setup loader" -> 1
         )
@@ -85,7 +86,6 @@ object LeakHygieneTests extends UtestIntegrationTestSuite {
           tester.eval(("show", "clean"))
           tester.eval(("show", "__.compile"))
           checkClassloaders(tester)(
-            "mill.codesig.ExternalSummary.apply upstreamClassloader" -> 1,
             "mill.daemon.MillBuildBootstrap#processRunClasspath classLoader cl" -> 1,
             "mill.kotlinlib.KotlinWorkerManager" -> 1,
             "mill.javalib.JvmWorkerModule#worker cl" -> 2,
@@ -110,7 +110,6 @@ object LeakHygieneTests extends UtestIntegrationTestSuite {
         for (i <- Range(0, 2)) {
           tester.eval(("show", "__.compile"))
           checkClassloaders(tester)(
-            "mill.codesig.ExternalSummary.apply upstreamClassloader" -> 1,
             "mill.daemon.MillBuildBootstrap#processRunClasspath classLoader cl" -> 1,
             "mill.kotlinlib.KotlinWorkerManager" -> 1,
             "mill.javalib.JvmWorkerModule#worker cl" -> 2,
@@ -133,7 +132,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite {
 
         // Exercise post-shutdown
 
-        tester.eval(("shutdown"))
+        tester.eval(("shutdown"), check = true)
         checkClassloaders(tester)(
           "mill.daemon.MillBuildBootstrap#processRunClasspath classLoader cl" -> 1,
           "mill.javalib.JvmWorkerModule#worker cl" -> 1

runner/codesig/src/mill/codesig/ExternalSummary.scala

@@ -6,6 +6,8 @@ import org.objectweb.asm.{ClassReader, ClassVisitor, MethodVisitor, Opcodes}
 
 import java.net.URLClassLoader
 
+import scala.util.Using
+
 case class ExternalSummary(
     directMethods: Map[JCls, Map[MethodSig, Boolean]],
     directAncestors: Map[JCls, Set[JCls]],
@@ -27,47 +29,49 @@ object ExternalSummary {
       localSummary: LocalSummary,
       upstreamClasspath: Seq[os.Path]
   )(implicit st: SymbolTable): ExternalSummary = {
-    val upstreamClassloader = mill.util.Jvm.createClassLoader(
+    def createUpstreamClassloader() = mill.util.Jvm.createClassLoader(
       upstreamClasspath,
       getClass.getClassLoader
     )
 
-    val allDirectAncestors = localSummary.mapValuesOnly(_.directAncestors).flatten
-
-    val allMethodCallParamClasses = localSummary
-      .mapValuesOnly(_.methods.values)
-      .flatten
-      .flatMap(_.calls)
-      .flatMap(call => Seq(call.cls) ++ call.desc.args)
-      .collect { case c: JType.Cls => c }
-
     val methodsPerCls = collection.mutable.Map.empty[JCls, Map[MethodSig, Boolean]]
     val ancestorsPerCls = collection.mutable.Map.empty[JCls, Set[JCls]]
     val directSuperclasses = collection.mutable.Map.empty[JCls, JCls]
 
-    def load(cls: JCls): Unit = methodsPerCls.getOrElse(cls, load0(cls))
-
-    def load0(cls: JCls): Unit = {
-      val visitor = new MyClassVisitor()
-      val resourcePath =
-        os.resource(upstreamClassloader) / os.SubPath(cls.name.replace('.', '/') + ".class")
-
-      new ClassReader(os.read.inputStream(resourcePath)).accept(
-        visitor,
-        ClassReader.SKIP_CODE | ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG
-      )
-
-      directSuperclasses(cls) = visitor.superclass
-      methodsPerCls(cls) = visitor.methods
-      ancestorsPerCls(cls) = visitor.ancestors
-      ancestorsPerCls(cls).foreach(load)
+    Using.resource(createUpstreamClassloader()) { upstreamClassloader =>
+      val allDirectAncestors = localSummary.mapValuesOnly(_.directAncestors).flatten
+
+      val allMethodCallParamClasses = localSummary
+        .mapValuesOnly(_.methods.values)
+        .flatten
+        .flatMap(_.calls)
+        .flatMap(call => Seq(call.cls) ++ call.desc.args)
+        .collect { case c: JType.Cls => c }
+
+      def load(cls: JCls): Unit = methodsPerCls.getOrElse(cls, load0(cls))
+
+      def load0(cls: JCls): Unit = {
+        val visitor = new MyClassVisitor()
+        val resourcePath =
+          os.resource(upstreamClassloader) / os.SubPath(cls.name.replace('.', '/') + ".class")
+
+        new ClassReader(os.read.inputStream(resourcePath)).accept(
+          visitor,
+          ClassReader.SKIP_CODE | ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG
+        )
+
+        directSuperclasses(cls) = visitor.superclass
+        methodsPerCls(cls) = visitor.methods
+        ancestorsPerCls(cls) = visitor.ancestors
+        ancestorsPerCls(cls).foreach(load)
+      }
+
+      (allDirectAncestors ++ allMethodCallParamClasses)
+        .filter(!localSummary.contains(_))
+        .toSet
+        .foreach(load)
     }
 
-    (allDirectAncestors ++ allMethodCallParamClasses)
-      .filter(!localSummary.contains(_))
-      .toSet
-      .foreach(load)
-
     ExternalSummary(methodsPerCls.toMap, ancestorsPerCls.toMap, directSuperclasses.toMap)
   }