From a6efead6bafd34a7c55a58e3dc5d6267345672f1 Mon Sep 17 00:00:00 2001
From: aosagie <osagie@gmail.com>
Date: Sun, 2 Sep 2018 02:21:35 -0400
Subject: [PATCH] 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
---
 main/core/src/mill/util/ClassLoader.scala       | 12 +-----------
 main/src/mill/modules/Jvm.scala                 | 17 ++++++++++-------
 scalalib/src/mill/scalalib/TestRunner.scala     |  3 ++-
 .../mill/scalanativelib/ScalaNativeModule.scala |  2 +-
 4 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/main/core/src/mill/util/ClassLoader.scala b/main/core/src/mill/util/ClassLoader.scala
index c0421a7be04..17268fb0890 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 96992d20fce..f0c0d8a902f 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()
+      }
     }
   }
 
diff --git a/scalalib/src/mill/scalalib/TestRunner.scala b/scalalib/src/mill/scalalib/TestRunner.scala
index 1ea819acbf4..947021ba393 100644
--- a/scalalib/src/mill/scalalib/TestRunner.scala
+++ b/scalalib/src/mill/scalalib/TestRunner.scala
@@ -68,7 +68,8 @@ object TestRunner {
                testClassfilePath: Agg[Path],
                args: Seq[String])
               (implicit ctx: Ctx.Log with Ctx.Home): (String, Seq[mill.scalalib.TestRunner.Result]) = {
-    Jvm.inprocess(entireClasspath, classLoaderOverrideSbtTesting = true, isolated = true, cl => {
+    //Leave the context class loader set and open so that shutdown hooks can access it
+    Jvm.inprocess(entireClasspath, classLoaderOverrideSbtTesting = true, isolated = true, closeContextClassLoaderWhenDone = false, cl => {
       val frameworks = frameworkInstances(cl)
 
       val events = mutable.Buffer.empty[Event]
diff --git a/scalanativelib/src/mill/scalanativelib/ScalaNativeModule.scala b/scalanativelib/src/mill/scalanativelib/ScalaNativeModule.scala
index 748afb79ea8..c8d9abda31e 100644
--- a/scalanativelib/src/mill/scalanativelib/ScalaNativeModule.scala
+++ b/scalanativelib/src/mill/scalanativelib/ScalaNativeModule.scala
@@ -265,7 +265,7 @@ trait TestScalaNativeModule extends ScalaNativeModule with TestModule { testOute
     val frameworkInstances = TestRunner.frameworks(testFrameworks()) _
 
     val testClasses =
-      Jvm.inprocess(testClasspathJvm().map(_.path), classLoaderOverrideSbtTesting = true, isolated = true,
+      Jvm.inprocess(testClasspathJvm().map(_.path), classLoaderOverrideSbtTesting = true, isolated = true, closeContextClassLoaderWhenDone = true,
         cl => {
           frameworkInstances(cl).flatMap { framework =>
             val df = Lib.discoverTests(cl, framework, Agg(compile().classes.path))