Manage classloader in BackgroundJobService

In https://github.com/sbt/sbt/issues/5075 we realized that sbt 1.3.0
introduces a regression where it closes the classloader used to invoke
the main method for in process run before all of its non-daemon threads
have terminated. To fix this and still close the classloader, I add a
method, runInBackgroundWithLoader that provides the background job
services with an optional classloader that it can close after the job
completes.

This cleanly merges and works with 1.3.x as well.

main/src/main/scala/sbt/BackgroundJobService.scala

@@ -26,6 +26,21 @@ abstract class BackgroundJobService extends Closeable {
       start: (Logger, File) => Unit
   ): JobHandle
 
+  /**
+   * Launch a background job which is a function that runs inside another thread.
+   *  Differs from [[BackgroundJobService.runInBackground]] in that the start task
+   *  provides both an optional classloader and the work to run in the background thread.
+   *  The service should close the classloader when the task completes.
+   *  killing the job will interrupt() the thread. If your thread blocks on a process,
+   *  then you should get an InterruptedException while blocking on the process, and
+   *  then you could process.destroy() for example.
+   */
+  private[sbt] def runInBackgroundWithLoader(spawningTask: ScopedKey[_], state: State)(
+      start: (Logger, File) => (Option[ClassLoader], () => Unit)
+  ): JobHandle = runInBackground(spawningTask, state) { (logger, file) =>
+    start(logger, file)._2.apply()
+  }
+
   /** Same as shutown. */
   def close(): Unit
 

main/src/main/scala/sbt/Defaults.scala

@@ -1434,12 +1434,19 @@ object Defaults extends BuildCommon {
       val service = bgJobService.value
       val (mainClass, args) = parser.parsed
       val hashClasspath = (bgHashClasspath in bgRunMain).value
-      service.runInBackground(resolvedScoped.value, state.value) { (logger, workingDir) =>
-        val cp =
+      service.runInBackgroundWithLoader(resolvedScoped.value, state.value) { (logger, workingDir) =>
+        val files =
           if (copyClasspath.value)
             service.copyClasspath(products.value, classpath.value, workingDir, hashClasspath)
           else classpath.value
-        scalaRun.value.run(mainClass, data(cp), args, logger).get
+        val cp = data(files)
+        scalaRun.value match {
+          case r: Run =>
+            val loader = r.newLoader(cp)
+            (Some(loader), () => r.runWithLoader(loader, cp, mainClass, args, logger).get)
+          case sr =>
+            (None, () => sr.run(mainClass, cp, args, logger).get)
+        }
       }
     }
   }
@@ -1457,12 +1464,20 @@ object Defaults extends BuildCommon {
       val service = bgJobService.value
       val mainClass = mainClassTask.value getOrElse sys.error("No main class detected.")
       val hashClasspath = (bgHashClasspath in bgRun).value
-      service.runInBackground(resolvedScoped.value, state.value) { (logger, workingDir) =>
-        val cp =
+      service.runInBackgroundWithLoader(resolvedScoped.value, state.value) { (logger, workingDir) =>
+        val files =
           if (copyClasspath.value)
             service.copyClasspath(products.value, classpath.value, workingDir, hashClasspath)
           else classpath.value
-        scalaRun.value.run(mainClass, data(cp), parser.parsed, logger).get
+        val cp = data(files)
+        val args = parser.parsed
+        scalaRun.value match {
+          case r: Run =>
+            val loader = r.newLoader(cp)
+            (Some(loader), () => r.runWithLoader(loader, cp, mainClass, args, logger).get)
+          case sr =>
+            (None, () => sr.run(mainClass, cp, args, logger).get)
+        }
       }
     }
   }

main/src/main/scala/sbt/internal/DefaultBackgroundJobService.scala

@@ -16,6 +16,7 @@ import java.util.concurrent.atomic.AtomicLong
 
 import sbt.Def.{ Classpath, ScopedKey, Setting }
 import sbt.Scope.GlobalScope
+import sbt.internal.inc.classpath.ClasspathFilter
 import sbt.internal.util.{ Attributed, ManagedLogger }
 import sbt.io.syntax._
 import sbt.io.{ Hash, IO }
@@ -148,6 +149,12 @@ private[sbt] abstract class AbstractBackgroundJobService extends BackgroundJobSe
     pool.run(this, spawningTask, state)(start)
   }
 
+  override private[sbt] def runInBackgroundWithLoader(spawningTask: ScopedKey[_], state: State)(
+      start: (Logger, File) => (Option[ClassLoader], () => Unit)
+  ): JobHandle = {
+    pool.runWithLoader(this, spawningTask, state)(start)
+  }
+
   override final def close(): Unit = shutdown()
   override def shutdown(): Unit = {
     while (jobSet.nonEmpty) {
@@ -419,6 +426,20 @@ private[sbt] class BackgroundThreadPool extends java.io.Closeable {
         }
       }
   }
+  private class BackgroundRunnableWithLoader(
+      val loader: Option[ClassLoader],
+      taskName: String,
+      body: () => Unit
+  ) extends BackgroundRunnable(taskName, body) {
+    override def awaitTermination(): Unit = {
+      try super.awaitTermination()
+      finally loader.foreach {
+        case ac: AutoCloseable   => ac.close()
+        case cp: ClasspathFilter => cp.close()
+        case _                   =>
+      }
+    }
+  }
 
   def run(manager: AbstractBackgroundJobService, spawningTask: ScopedKey[_], state: State)(
       work: (Logger, File) => Unit
@@ -433,6 +454,22 @@ private[sbt] class BackgroundThreadPool extends java.io.Closeable {
     manager.doRunInBackground(spawningTask, state, start _)
   }
 
+  private[sbt] def runWithLoader(
+      manager: AbstractBackgroundJobService,
+      spawningTask: ScopedKey[_],
+      state: State
+  )(
+      getWork: (Logger, File) => (Option[ClassLoader], () => Unit)
+  ): JobHandle = {
+    def start(logger: Logger, workingDir: File): BackgroundJob = {
+      val (loader, work) = getWork(logger, workingDir)
+      val runnable = new BackgroundRunnableWithLoader(loader, spawningTask.key.label, work)
+      executor.execute(runnable)
+      runnable
+    }
+    manager.doRunInBackground(spawningTask, state, start _)
+  }
+
   override def close(): Unit = {
     executor.shutdown()
   }

run/src/main/scala/sbt/Run.scala

@@ -10,10 +10,9 @@ package sbt
 import java.io.File
 import java.lang.reflect.Method
 import java.lang.reflect.Modifier.{ isPublic, isStatic }
-import java.net.URLClassLoader
 
 import sbt.internal.inc.ScalaInstance
-import sbt.internal.inc.classpath.{ ClasspathFilter, ClasspathUtilities }
+import sbt.internal.inc.classpath.ClasspathUtilities
 import sbt.internal.util.MessageOnlyException
 import sbt.io.Path
 import sbt.util.Logger
@@ -59,17 +58,25 @@ class ForkRun(config: ForkOptions) extends ScalaRun {
   private def classpathOption(classpath: Seq[File]) =
     "-classpath" :: Path.makeString(classpath) :: Nil
 }
-class Run(newLoader: Seq[File] => ClassLoader, trapExit: Boolean) extends ScalaRun {
+class Run(private[sbt] val newLoader: Seq[File] => ClassLoader, trapExit: Boolean)
+    extends ScalaRun {
   def this(instance: ScalaInstance, trapExit: Boolean, nativeTmp: File) =
     this((cp: Seq[File]) => ClasspathUtilities.makeLoader(cp, instance, nativeTmp), trapExit)
 
-  /** Runs the class 'mainClass' using the given classpath and options using the scala runner.*/
-  def run(mainClass: String, classpath: Seq[File], options: Seq[String], log: Logger): Try[Unit] = {
+  private[sbt] def runWithLoader(
+      loader: ClassLoader,
+      classpath: Seq[File],
+      mainClass: String,
+      options: Seq[String],
+      log: Logger
+  ): Try[Unit] = {
     log.info(s"running $mainClass ${Run.runOptionsStr(options)}")
 
-    def execute() =
+    def execute(): Unit =
       try {
-        run0(mainClass, classpath, options, log)
+        log.debug("  Classpath:\n\t" + classpath.mkString("\n\t"))
+        val main = getMainMethod(mainClass, loader)
+        invokeMain(loader, main, options)
       } catch {
         case e: java.lang.reflect.InvocationTargetException => throw e.getCause
       }
@@ -85,24 +92,10 @@ class Run(newLoader: Seq[File] => ClassLoader, trapExit: Boolean) extends ScalaR
     if (trapExit) Run.executeTrapExit(execute(), log)
     else directExecute()
   }
-  private def run0(
-      mainClassName: String,
-      classpath: Seq[File],
-      options: Seq[String],
-      log: Logger
-  ): Unit = {
-    log.debug("  Classpath:\n\t" + classpath.mkString("\n\t"))
-    val loader = newLoader(classpath)
-    try {
-      val main = getMainMethod(mainClassName, loader)
-      invokeMain(loader, main, options)
-    } finally {
-      loader match {
-        case u: URLClassLoader  => u.close()
-        case c: ClasspathFilter => c.close()
-        case _                  =>
-      }
-    }
+
+  /** Runs the class 'mainClass' using the given classpath and options using the scala runner.*/
+  def run(mainClass: String, classpath: Seq[File], options: Seq[String], log: Logger): Try[Unit] = {
+    runWithLoader(newLoader(classpath), classpath, mainClass, options, log)
   }
   private def invokeMain(
       loader: ClassLoader,