Made scripts wait for all non-daemon threads to...

Made scripts wait for all non-daemon threads to exit before calling
System.exit. Closes #1955, #2006, #3408. Review by community.

3 files changed, 51 insertions, 23 deletions

diff --git a/src/compiler/scala/tools/nsc/MainGenericRunner.scala b/src/compiler/scala/tools/nsc/MainGenericRunner.scala
index 936ee3c1db..96d6846d64 100644
--- a/src/compiler/scala/tools/nsc/MainGenericRunner.scala
+++ b/src/compiler/scala/tools/nsc/MainGenericRunner.scala
@@ -13,7 +13,7 @@ import java.net.{ URL, MalformedURLException }
 import scala.tools.util.PathResolver
 
 import io.{ File, Process }
-import util.{ ClassPath, ScalaClassLoader }
+import util.{ ClassPath, ScalaClassLoader, waitingForThreads }
 import Properties.{ versionString, copyrightString }
 
 /** An object that runs Scala code.  It has three possible
diff --git a/src/compiler/scala/tools/nsc/ScriptRunner.scala b/src/compiler/scala/tools/nsc/ScriptRunner.scala
index a7d67a3af9..088efe40ff 100644
--- a/src/compiler/scala/tools/nsc/ScriptRunner.scala
+++ b/src/compiler/scala/tools/nsc/ScriptRunner.scala
@@ -17,6 +17,7 @@ import java.lang.reflect.InvocationTargetException
 import java.net.URL
 import java.util.jar.{ JarEntry, JarOutputStream }
 
+import util.waitingForThreads
 import scala.tools.util.PathResolver
 import scala.tools.nsc.reporters.{Reporter,ConsoleReporter}
 
@@ -198,31 +199,36 @@ object ScriptRunner {
       else None
     }
 
-    if (settings.savecompiled.value) {
-      val jarFile = jarFileFor(scriptFile)
-      def jarOK   = jarFile.canRead && (jarFile isFresher File(scriptFile))
-
-      def recompile() = {
-        jarFile.delete()
-
-        compile match {
-          case Some(compiledPath) =>
-            tryMakeJar(jarFile, compiledPath)
-            if (jarOK) {
-              compiledPath.deleteRecursively()
-              handler(jarFile.toAbsolute.path)
-            }
-            // jar failed; run directly from the class files
-            else handler(compiledPath.path)
-          case _  => false
+    /** The script runner calls System.exit to communicate a return value, but this must
+     *  not take place until there are no non-daemon threads running.  Tickets #1955, #2006.
+     */
+    waitingForThreads {
+      if (settings.savecompiled.value) {
+        val jarFile = jarFileFor(scriptFile)
+        def jarOK   = jarFile.canRead && (jarFile isFresher File(scriptFile))
+
+        def recompile() = {
+          jarFile.delete()
+
+          compile match {
+            case Some(compiledPath) =>
+              tryMakeJar(jarFile, compiledPath)
+              if (jarOK) {
+                compiledPath.deleteRecursively()
+                handler(jarFile.toAbsolute.path)
+              }
+              // jar failed; run directly from the class files
+              else handler(compiledPath.path)
+            case _  => false
+          }
         }
-      }
 
-      if (jarOK) handler(jarFile.toAbsolute.path) // pre-compiled jar is current
-      else recompile()                            // jar old - recompile the script.
+        if (jarOK) handler(jarFile.toAbsolute.path) // pre-compiled jar is current
+        else recompile()                            // jar old - recompile the script.
+      }
+      // don't use a cache jar at all--just use the class files
+      else compile map (cp => handler(cp.path)) getOrElse false
     }
-    // don't use a cache jar at all--just use the class files
-    else compile map (cp => handler(cp.path)) getOrElse false
   }
 
   /** Run a script after it has been compiled
diff --git a/src/compiler/scala/tools/nsc/util/package.scala b/src/compiler/scala/tools/nsc/util/package.scala
index 92d4eab54f..c38b2c5031 100644
--- a/src/compiler/scala/tools/nsc/util/package.scala
+++ b/src/compiler/scala/tools/nsc/util/package.scala
@@ -11,6 +11,28 @@ package object util {
   /** Apply a function and return the passed value */
   def returning[T](x: T)(f: T => Unit): T = { f(x) ; x }
 
+  /** All living threads. */
+  def allThreads(): List[Thread] = {
+    val num = Thread.activeCount()
+    val tarray = new Array[Thread](num)
+    val got = Thread.enumerate(tarray)
+
+    tarray take got toList
+  }
+
+  /** Execute code and then wait for all Threads created during its
+   *  execution to complete.
+   */
+  def waitingForThreads[T](body: => T) = {
+    val ts1 = allThreads()
+    val result = body
+    val ts2 = allThreads()
+    val newThreads = (ts2.toSet -- ts1) filterNot (_.isDaemon())
+
+    newThreads foreach (_.join())
+    result
+  }
+
   /** Generate a string using a routine that wants to write on a stream. */
   def stringFromWriter(writer: PrintWriter => Unit): String = {
     val stringWriter = new StringWriter()