Manage shutdown hooks

I discovered that some registered shutdown hooks would crash due to
67df72ab01 because they would try to load
classes from the closed classloader. To fix this, I add a internal
shutdown hooks mechanism that can be managed by sbt. Any unevaluated
shutdown hooks will be run when the sbt main method exits. This means
that they will be run when the user calls reboot. I think that is
reasonable.

build.sbt

@@ -666,6 +666,7 @@ lazy val mainProj = (project in file("main"))
       exclude[DirectMissingMethodProblem]("sbt.Defaults.allTestGroupsTask"),
       exclude[DirectMissingMethodProblem]("sbt.Plugins.topologicalSort"),
       exclude[IncompatibleMethTypeProblem]("sbt.Defaults.allTestGroupsTask"),
+      exclude[DirectMissingMethodProblem]("sbt.StandardMain.shutdownHook")
     )
   )
   .configure(

main/src/main/scala/sbt/Main.scala

@@ -98,33 +98,36 @@ final class xMain extends xsbti.AppMain {
   }
 }
 private[sbt] object xMainImpl {
-  private[sbt] def run(configuration: xsbti.AppConfiguration): xsbti.MainResult = {
+  private[sbt] def run(configuration: xsbti.AppConfiguration): xsbti.MainResult =
-    import BasicCommandStrings.{ DashClient, DashDashClient, runEarly }
+    try {
-    import BasicCommands.early
+      import BasicCommandStrings.{ DashClient, DashDashClient, runEarly }
-    import BuiltinCommands.defaults
+      import BasicCommands.early
-    import sbt.internal.CommandStrings.{ BootCommand, DefaultsCommand, InitCommand }
+      import BuiltinCommands.defaults
-    import sbt.internal.client.NetworkClient
+      import sbt.internal.CommandStrings.{ BootCommand, DefaultsCommand, InitCommand }
+      import sbt.internal.client.NetworkClient
 
-    // if we detect -Dsbt.client=true or -client, run thin client.
+      // if we detect -Dsbt.client=true or -client, run thin client.
-    val clientModByEnv = java.lang.Boolean.getBoolean("sbt.client")
+      val clientModByEnv = java.lang.Boolean.getBoolean("sbt.client")
-    val userCommands = configuration.arguments.map(_.trim)
+      val userCommands = configuration.arguments.map(_.trim)
-    if (clientModByEnv || (userCommands.exists { cmd =>
+      if (clientModByEnv || (userCommands.exists { cmd =>
+            (cmd == DashClient) || (cmd == DashDashClient)
+          })) {
+        val args = userCommands.toList filterNot { cmd =>
           (cmd == DashClient) || (cmd == DashDashClient)
-        })) {
+        }
-      val args = userCommands.toList filterNot { cmd =>
+        NetworkClient.run(configuration, args)
-        (cmd == DashClient) || (cmd == DashDashClient)
+        Exit(0)
+      } else {
+        val state = StandardMain.initialState(
+          configuration,
+          Seq(defaults, early),
+          runEarly(DefaultsCommand) :: runEarly(InitCommand) :: BootCommand :: Nil
+        )
+        StandardMain.runManaged(state)
       }
-      NetworkClient.run(configuration, args)
+    } finally {
-      Exit(0)
+      ShutdownHooks.close()
-    } else {
-      val state = StandardMain.initialState(
-        configuration,
-        Seq(defaults, early),
-        runEarly(DefaultsCommand) :: runEarly(InitCommand) :: BootCommand :: Nil
-      )
-      StandardMain.runManaged(state)
     }
-  }
 }
 
 final class ScriptMain extends xsbti.AppMain {
@@ -155,30 +158,21 @@ object StandardMain {
   import scalacache.caffeine._
   private[sbt] lazy val cache: scalacache.Cache[Any] = CaffeineCache[Any]
 
-  private[this] val closeRunnable: Runnable = () => {
+  private[this] val closeRunnable = () => {
     cache.close()(scalacache.modes.sync.mode)
     cache.close()(scalacache.modes.scalaFuture.mode(ExecutionContext.global))
     exchange.shutdown()
   }
-  private[sbt] val shutdownHook = new Thread(closeRunnable)
 
   def runManaged(s: State): xsbti.MainResult = {
     val previous = TrapExit.installManager()
     try {
       try {
-        val hooked = try {
+        val hook = ShutdownHooks.add(closeRunnable)
-          Runtime.getRuntime.addShutdownHook(shutdownHook)
-          true
-        } catch {
-          case _: IllegalArgumentException => false
-        }
         try {
           MainLoop.runLogged(s)
         } finally {
-          closeRunnable.run()
+          hook.close()
-          if (hooked) {
-            Runtime.getRuntime.removeShutdownHook(shutdownHook)
-          }
           ()
         }
       } finally DefaultBackgroundJobService.backgroundJobService.shutdown()

main/src/main/scala/sbt/MainLoop.scala

@@ -11,6 +11,7 @@ import java.io.PrintWriter
 import java.util.Properties
 
 import jline.TerminalFactory
+import sbt.internal.ShutdownHooks
 import sbt.internal.langserver.ErrorCodes
 import sbt.internal.util.{ ErrorHandling, GlobalLogBacking }
 import sbt.io.{ IO, Using }
@@ -27,13 +28,12 @@ object MainLoop {
     // We've disabled jline shutdown hooks to prevent classloader leaks, and have been careful to always restore
     // the jline terminal in finally blocks, but hitting ctrl+c prevents finally blocks from being executed, in that
     // case the only way to restore the terminal is in a shutdown hook.
-    val shutdownHook = new Thread(() => TerminalFactory.get().restore())
+    val shutdownHook = ShutdownHooks.add(() => TerminalFactory.get().restore())
 
     try {
-      Runtime.getRuntime.addShutdownHook(shutdownHook)
       runLoggedLoop(state, state.globalLogging.backing)
     } finally {
-      Runtime.getRuntime.removeShutdownHook(shutdownHook)
+      shutdownHook.close()
       ()
     }
   }

main/src/main/scala/sbt/internal/LayeredClassLoader.scala

@@ -49,12 +49,10 @@ private[sbt] class LayeredClassLoader(
 
 private[internal] object NativeLibs {
   private[this] val nativeLibs = new jutil.HashSet[File].asScala
-  Runtime.getRuntime.addShutdownHook(new Thread("sbt.internal.native-library-deletion") {
+  ShutdownHooks.add(() => {
-    override def run(): Unit = {
+    nativeLibs.foreach(IO.delete)
-      nativeLibs.foreach(IO.delete)
+    IO.deleteIfEmpty(nativeLibs.map(_.getParentFile).toSet)
-      IO.deleteIfEmpty(nativeLibs.map(_.getParentFile).toSet)
+    nativeLibs.clear()
-      nativeLibs.clear()
-    }
   })
   def addNativeLib(lib: String): Unit = {
     nativeLibs.add(new File(lib))

main/src/main/scala/sbt/internal/ShutdownHooks.scala

@@ -0,0 +1,45 @@
+/*
+ * sbt
+ * Copyright 2011 - 2018, Lightbend, Inc.
+ * Copyright 2008 - 2010, Mark Harrah
+ * Licensed under Apache License 2.0 (see LICENSE)
+ */
+
+package sbt.internal
+
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.{ AtomicBoolean, AtomicInteger }
+
+import scala.util.control.NonFatal
+
+private[sbt] object ShutdownHooks extends AutoCloseable {
+  private[this] val idGenerator = new AtomicInteger(0)
+  private[this] val hooks = new ConcurrentHashMap[Int, () => Any]
+  private[this] val ranHooks = new AtomicBoolean(false)
+  private[this] val thread = new Thread("shutdown-hooks-run-all") {
+    override def run(): Unit = runAll()
+  }
+  private[this] val runtime = Runtime.getRuntime
+  runtime.addShutdownHook(thread)
+  private[sbt] def add[R](task: () => R): AutoCloseable = {
+    val id = idGenerator.getAndIncrement()
+    hooks.put(
+      id,
+      () =>
+        try task()
+        catch {
+          case NonFatal(e) =>
+            System.err.println(s"Caught exception running shutdown hook: $e")
+            e.printStackTrace(System.err)
+        }
+    )
+    () => Option(hooks.remove(id)).foreach(_.apply())
+  }
+  private def runAll(): Unit = if (ranHooks.compareAndSet(false, true)) {
+    hooks.forEachValue(Runtime.getRuntime.availableProcessors, _.apply())
+  }
+  override def close(): Unit = {
+    runtime.removeShutdownHook(thread)
+    runAll()
+  }
+}

main/src/main/scala/sbt/internal/TaskTimings.scala

@@ -38,9 +38,7 @@ private[sbt] final class TaskTimings(reportOnShutdown: Boolean)
 
   if (reportOnShutdown) {
     start = System.nanoTime
-    Runtime.getRuntime.addShutdownHook(new Thread {
+    ShutdownHooks.add(() => report())
-      override def run() = report()
-    })
   }
 
   override def initial(): Unit = {

main/src/main/scala/sbt/internal/TaskTraceEvent.scala

@@ -36,9 +36,7 @@ private[sbt] final class TaskTraceEvent
   override def stop(): Unit = ()
 
   start = System.nanoTime
-  Runtime.getRuntime.addShutdownHook(new Thread {
+  ShutdownHooks.add(() => report())
-    override def run() = report()
-  })
 
   private[this] def report() = {
     if (timings.asScala.nonEmpty) {