From be489e05caeb4db85d7aa5e3a36f0bf340f74cc7 Mon Sep 17 00:00:00 2001
From: Ethan Atkins <contact@ethanatkins.com>
Date: Sun, 28 Jul 2019 15:08:22 -0700
Subject: [PATCH] Clear expired loaders

Sometimes turbo mode didn't work correctly for projects where resources
were modified. This was because it was possible for the resource
classloader to inadvertently evict the dependency classloader from the
classloader cache because they had the same file stamps. There were two
fixes:
1) remove expired entries from the cache based on the
    (Parent, Classpath) pair rather than just classpath
2) do not close the classloaders during cache eviction. They may still
   be in use when we evict them so we need to wait until they are
   explicitly closed elsewhere or until the go out of scope and are
   collected by the CleanupThread

I tested this change with a spark project in which I kept modifying the
resources. Prior to this change, I could get into a state where if I
modified the resources, the dependency layer would get evicted every
time so the benefits of turbo mode were not realized.
---
 .../sbt/internal/classpath/WrappedLoader.java | 10 ++++++
 .../internal/classpath/ClassLoaderCache.scala | 34 ++++++++++++-------
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/main-command/src/main/java/sbt/internal/classpath/WrappedLoader.java b/main-command/src/main/java/sbt/internal/classpath/WrappedLoader.java
index 8bf7b7edc..932187f65 100644
--- a/main-command/src/main/java/sbt/internal/classpath/WrappedLoader.java
+++ b/main-command/src/main/java/sbt/internal/classpath/WrappedLoader.java
@@ -9,8 +9,10 @@ package sbt.internal.classpath;
 
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class WrappedLoader extends URLClassLoader {
+  private final AtomicBoolean invalidated = new AtomicBoolean(false);
   static {
     ClassLoader.registerAsParallelCapable();
   }
@@ -19,6 +21,14 @@ public class WrappedLoader extends URLClassLoader {
     super(new URL[] {}, parent);
   }
 
+  void invalidate() {
+    invalidated.set(true);
+  }
+
+  boolean invalidated() {
+    return invalidated.get();
+  }
+
   @Override
   public URL[] getURLs() {
     final ClassLoader parent = getParent();
diff --git a/main-command/src/main/scala/sbt/internal/classpath/ClassLoaderCache.scala b/main-command/src/main/scala/sbt/internal/classpath/ClassLoaderCache.scala
index c0683d040..eaf3ff1de 100644
--- a/main-command/src/main/scala/sbt/internal/classpath/ClassLoaderCache.scala
+++ b/main-command/src/main/scala/sbt/internal/classpath/ClassLoaderCache.scala
@@ -71,18 +71,26 @@ private[sbt] class ClassLoaderCache(
     new java.util.concurrent.ConcurrentHashMap[Key, Reference[ClassLoader]]()
   private[this] val referenceQueue = new ReferenceQueue[ClassLoader]
 
-  private[this] def closeExpiredLoaders(): Unit = {
-    val toClose = lock.synchronized(delegate.asScala.groupBy(_._1.files.toSet).flatMap {
+  private[this] def clearExpiredLoaders(): Unit = lock.synchronized {
+    val clear = (k: Key, ref: Reference[ClassLoader]) => {
+      ref.get() match {
+        case w: WrappedLoader => w.invalidate()
+        case _                =>
+      }
+      delegate.remove(k)
+      ()
+    }
+    def isInvalidated(classLoader: ClassLoader): Boolean = classLoader match {
+      case w: WrappedLoader => w.invalidated()
+      case _                => false
+    }
+    delegate.asScala.groupBy { case (k, _) => k.parent -> k.files.toSet }.foreach {
       case (_, pairs) if pairs.size > 1 =>
-        val max = pairs.maxBy(_._1.maxStamp)._1
-        pairs.filterNot(_._1 == max).flatMap {
-          case (k, v) =>
-            delegate.remove(k)
-            Option(v.get)
-        }
-      case _ => Nil
-    })
-    toClose.foreach(close)
+        val max = pairs.map(_._1.maxStamp).max
+        pairs.foreach { case (k, v) => if (k.maxStamp != max) clear(k, v) }
+      case _ =>
+    }
+    delegate.forEach((k, v) => if (isInvalidated(k.parent)) clear(k, v))
   }
   private[this] class CleanupThread(private[this] val id: Int)
       extends Thread(s"classloader-cache-cleanup-$id") {
@@ -97,7 +105,7 @@ private[sbt] class ClassLoaderCache(
             delegate.remove(key)
           case _ =>
         }
-        closeExpiredLoaders()
+        clearExpiredLoaders()
         false
       } catch {
         case _: InterruptedException => true
@@ -178,7 +186,7 @@ private[sbt] class ClassLoaderCache(
           val ref = mkReference(key, f())
           val loader = ref.get
           delegate.put(key, ref)
-          closeExpiredLoaders()
+          clearExpiredLoaders()
           loader
         }
         lock.synchronized {