Rework closure of macro classloaders

Author: retronym

src/compiler/scala/reflect/macros/runtime/MacroRuntimes.scala

@@ -54,23 +54,8 @@ trait MacroRuntimes extends JavaReflectionRuntimes {
   /** Macro classloader that is used to resolve and run macro implementations.
    *  Loads classes from from -cp (aka the library classpath).
    *  Is also capable of detecting REPL and reusing its classloader.
-   *
-   *  When -Xmacro-jit is enabled, we sometimes fallback to on-the-fly compilation of macro implementations,
-   *  which compiles implementations into a virtual directory (very much like REPL does) and then conjures
-   *  a classloader mapped to that virtual directory.
    */
-  private lazy val defaultMacroClassloaderCache = {
-    def attemptClose(loader: ClassLoader): Unit = {
-      if (!scala.tools.nsc.typechecker.Macros.macroClassLoadersCache.owns(loader)) {
-        loader match {
-          case u: URLClassLoader => debuglog("Closing macro runtime classloader"); u.close()
-          case afcl: AbstractFileClassLoader => attemptClose(afcl.getParent)
-          case _ => ???
-        }
-      }
-    }
-    perRunCaches.newGeneric(findMacroClassLoader, attemptClose _)
-  }
+  private lazy val defaultMacroClassloaderCache: () => ClassLoader = perRunCaches.newGeneric(findMacroClassLoader())
   def defaultMacroClassloader: ClassLoader = defaultMacroClassloaderCache()
 
   /** Abstracts away resolution of macro runtimes.

src/compiler/scala/tools/nsc/classpath/ZipAndJarFileLookupFactory.scala

@@ -190,7 +190,6 @@ final class FileBasedCache[T] {
   import java.nio.file.Path
   private case class Stamp(lastModified: FileTime, fileKey: Object)
   private val cache = collection.mutable.Map.empty[Seq[Path], (Seq[Stamp], T)]
-  def owns(t: T): Boolean = cache.valuesIterator.exists(_._2.asInstanceOf[AnyRef] eq t.asInstanceOf[AnyRef])
 
   def getOrCreate(paths: Seq[Path], create: () => T): T = cache.synchronized {
     val stamps = paths.map { path =>

src/compiler/scala/tools/nsc/plugins/Plugins.scala

@@ -164,13 +164,13 @@ trait Plugins { global: Global =>
         // TODO if the only null is jrt:// we can still cache
         // TODO filter out classpath elements pointing to non-existing files before we get here, that's another source of null
         analyzer.macroLogVerbose(s"macro classloader: caching is disabled because `AbstractFile.getURL` returned `null` for ${hasNullURL.map(_._1).mkString(", ")}.")
-        newLoader()
+        perRunCaches.recordClassloader(newLoader())
       } else {
         val locations = urlsAndFiles.map(t => Path(t._2.file))
         val nonJarZips = locations.filterNot(Jar.isJarOrZip(_))
         if (nonJarZips.nonEmpty) {
           analyzer.macroLogVerbose(s"macro classloader: caching is disabled because the following paths are not supported: ${nonJarZips.mkString(",")}.")
-          newLoader()
+          perRunCaches.recordClassloader(newLoader())
         } else {
           Macros.macroClassLoadersCache.getOrCreate(locations.map(_.jfile.toPath()), newLoader)
         }

src/compiler/scala/tools/reflect/ReflectGlobal.scala

@@ -32,7 +32,7 @@ class ReflectGlobal(currentSettings: Settings, reporter: Reporter, override val
     */
   override protected[scala] def findMacroClassLoader(): ClassLoader = {
     val classpath = classPath.asURLs
-    ScalaClassLoader.fromURLs(classpath, rootClassLoader)
+    perRunCaches.recordClassloader(ScalaClassLoader.fromURLs(classpath, rootClassLoader))
   }
 
   override def transformedType(sym: Symbol) =

src/reflect/scala/reflect/internal/SymbolTable.scala

@@ -14,6 +14,8 @@ package scala
 package reflect
 package internal
 
+import java.net.URLClassLoader
+
 import scala.annotation.elidable
 import scala.collection.mutable
 import util._
@@ -415,6 +417,22 @@ abstract class SymbolTable extends macros.Universe
       cache
     }
 
+    /** Closes the provided classloader at the conclusion of this Run */
+    final def recordClassloader(loader: ClassLoader): ClassLoader = {
+      def attemptClose(loader: ClassLoader): Unit = {
+        loader match {
+          case u: URLClassLoader => debuglog("Closing classloader " + u); u.close()
+          case _ =>
+        }
+      }
+      caches ::= new WeakReference((new Clearable {
+        def clear(): Unit = {
+          attemptClose(loader)
+        }
+      }))
+      loader
+    }
+
     /**
      * Removes a cache from the per-run caches. This is useful for testing: it allows running the
      * compiler and then inspect the state of a cache.