Fix shutdown hook failures in tests (#422)

* Remove duplication from ClassLoader.create

* Prevent closing of context class loader in tests so that shutdown hooks can run

2 files changed, 11 insertions, 18 deletions

diff --git a/main/core/src/mill/util/ClassLoader.scala b/main/core/src/mill/util/ClassLoader.scala
index c0421a7b..17268fb0 100644
--- a/main/core/src/mill/util/ClassLoader.scala
+++ b/main/core/src/mill/util/ClassLoader.scala
@@ -8,17 +8,7 @@ import io.github.retronym.java9rtexport.Export
 object ClassLoader {
 
   def create(urls: Seq[URL], parent: java.lang.ClassLoader)(
-      implicit ctx: Ctx.Home): URLClassLoader = {
-    new URLClassLoader(
-      makeUrls(urls).toArray,
-      refinePlatformParent(parent)
-    ) {
-      override def findClass(name: String): Class[_] = {
-        if (name.startsWith("com.sun.jna")) getClass.getClassLoader.loadClass(name)
-        else super.findClass(name)
-      }
-    }
-  }
+      implicit ctx: Ctx.Home): URLClassLoader = create(urls, parent, _ => None)
 
   def create(urls: Seq[URL],
              parent: java.lang.ClassLoader,
diff --git a/main/src/mill/modules/Jvm.scala b/main/src/mill/modules/Jvm.scala
index 96992d20..f0c0d8a9 100644
--- a/main/src/mill/modules/Jvm.scala
+++ b/main/src/mill/modules/Jvm.scala
@@ -89,7 +89,7 @@ object Jvm {
                classPath: Agg[Path],
                mainArgs: Seq[String] = Seq.empty)
               (implicit ctx: Ctx): Unit = {
-    inprocess(classPath, classLoaderOverrideSbtTesting = false, isolated = true, cl => {
+    inprocess(classPath, classLoaderOverrideSbtTesting = false, isolated = true, closeContextClassLoaderWhenDone = true, cl => {
       getMainMethod(mainClass, cl).invoke(null, mainArgs.toArray)
     })
   }
@@ -113,6 +113,7 @@ object Jvm {
   def inprocess[T](classPath: Agg[Path],
                    classLoaderOverrideSbtTesting: Boolean,
                    isolated: Boolean,
+                   closeContextClassLoaderWhenDone: Boolean,
                    body: ClassLoader => T)
                   (implicit ctx: Ctx.Home): T = {
     val urls = classPath.map(_.toIO.toURI.toURL)
@@ -123,19 +124,21 @@ object Jvm {
           Some(outerClassLoader.loadClass(name))
         else None
       })
-    } else if (isolated){
-
+    } else if (isolated) {
       mill.util.ClassLoader.create(urls.toVector, null)
-    }else{
+    } else {
       mill.util.ClassLoader.create(urls.toVector, getClass.getClassLoader)
     }
+
     val oldCl = Thread.currentThread().getContextClassLoader
     Thread.currentThread().setContextClassLoader(cl)
     try {
       body(cl)
-    }finally{
-      Thread.currentThread().setContextClassLoader(oldCl)
-      cl.close()
+    } finally {
+      if (closeContextClassLoaderWhenDone) {
+        Thread.currentThread().setContextClassLoader(oldCl)
+        cl.close()
+      }
     }
   }