Join worker threads on pool shutdown

benchmarks/src/main/scala/cats/effect/benchmarks/WorkStealingBenchmark.scala

@@ -171,6 +171,7 @@ class WorkStealingBenchmark {
       "io-blocker",
       60.seconds,
       false,
+      1.second,
       SleepSystem,
       _.printStackTrace())
 

build.sbt

@@ -898,7 +898,8 @@ lazy val tests: CrossProject = crossProject(JSPlatform, JVMPlatform, NativePlatf
     scalacOptions ~= { _.filterNot(_.startsWith("-P:scalajs:mapSourceURI")) }
   )
   .jvmSettings(
-    Test / fork := true
+    fork := true,
+    Test / javaOptions += s"-Dsbt.classpath=${(Test / fullClasspath).value.map(_.data.getAbsolutePath).mkString(File.pathSeparator)}"
   )
   .nativeSettings(
     Compile / mainClass := Some("catseffect.examples.NativeRunner")

core/jvm/src/main/scala/cats/effect/unsafe/IORuntimeCompanionPlatform.scala

@@ -65,6 +65,7 @@ private[unsafe] abstract class IORuntimeCompanionPlatform { this: IORuntime.type
       runtimeBlockingExpiration,
       reportFailure,
       false,
+      1.second,
       SleepSystem
     )
     (pool, shutdown)
@@ -77,6 +78,7 @@ private[unsafe] abstract class IORuntimeCompanionPlatform { this: IORuntime.type
       runtimeBlockingExpiration: Duration = 60.seconds,
       reportFailure: Throwable => Unit = _.printStackTrace(),
       blockedThreadDetectionEnabled: Boolean = false,
+      shutdownTimeout: Duration = 1.second,
       pollingSystem: PollingSystem = SelectorSystem())
       : (WorkStealingThreadPool[_], pollingSystem.Api, () => Unit) = {
     val threadPool =
@@ -86,8 +88,10 @@ private[unsafe] abstract class IORuntimeCompanionPlatform { this: IORuntime.type
         blockerThreadPrefix,
         runtimeBlockingExpiration,
         blockedThreadDetectionEnabled && (threads > 1),
+        shutdownTimeout,
         pollingSystem,
-        reportFailure)
+        reportFailure
+      )
 
     val unregisterMBeans =
       if (isStackTracing) {

core/jvm/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala

@@ -64,6 +64,7 @@ private[effect] final class WorkStealingThreadPool[P](
     private[unsafe] val blockerThreadPrefix: String, // prefix for the name of worker threads currently in a blocking region
     private[unsafe] val runtimeBlockingExpiration: Duration,
     private[unsafe] val blockedThreadDetectionEnabled: Boolean,
+    shutdownTimeout: Duration,
     system: PollingSystem.WithPoller[P],
     reportFailure0: Throwable => Unit
 ) extends ExecutionContextExecutor
@@ -691,39 +692,72 @@ private[effect] final class WorkStealingThreadPool[P](
   def shutdown(): Unit = {
     // Clear the interrupt flag.
     val interruptCalling = Thread.interrupted()
+    val currentThread = Thread.currentThread()
 
     // Execute the shutdown logic only once.
     if (done.compareAndSet(false, true)) {
-      // Send an interrupt signal to each of the worker threads.
-      workerThreadPublisher.get()
-
       // Note: while loops and mutable variables are used throughout this method
       // to avoid allocations of objects, since this method is expected to be
       // executed mostly in situations where the thread pool is shutting down in
       // the face of unhandled exceptions or as part of the whole JVM exiting.
+
+      workerThreadPublisher.get()
+
+      // Send an interrupt signal to each of the worker threads.
       var i = 0
       while (i < threadCount) {
-        workerThreads(i).interrupt()
-        system.closePoller(pollers(i))
+        val workerThread = workerThreads(i)
+        if (workerThread ne currentThread) {
+          workerThread.interrupt()
+        }
         i += 1
       }
 
-      system.close()
+      i = 0
+      var joinTimeout = shutdownTimeout match {
+        case Duration.Inf => Long.MaxValue
+        case d => d.toNanos
+      }
+      while (i < threadCount && joinTimeout > 0) {
+        val workerThread = workerThreads(i)
+        if (workerThread ne currentThread) {
+          val now = System.nanoTime()
+          workerThread.join(joinTimeout / 1000000, (joinTimeout % 1000000).toInt)
+          val elapsed = System.nanoTime() - now
+          joinTimeout -= elapsed
+        }
+        i += 1
+      }
 
-      // Clear the interrupt flag.
-      Thread.interrupted()
+      i = 0
+      var allClosed = true
+      while (i < threadCount) {
+        val workerThread = workerThreads(i)
+        // only close the poller if it is safe to do so, leak otherwise ...
+        if ((workerThread eq currentThread) || !workerThread.isAlive()) {
+          system.closePoller(pollers(i))
+        } else {
+          allClosed = false
+        }
+        i += 1
+      }
+
+      if (allClosed) {
+        system.close()
+      }
 
       var t: WorkerThread[P] = null
       while ({
         t = cachedThreads.pollFirst()
         t ne null
       }) {
         t.interrupt()
+        // don't bother joining, cached threads are not doing anything interesting
       }
 
       // Drain the external queue.
       externalQueue.clear()
-      if (interruptCalling) Thread.currentThread().interrupt()
+      if (interruptCalling) currentThread.interrupt()
     }
   }
 

tests/jvm/src/main/scala/catseffect/examplesplatform.scala

@@ -36,7 +36,7 @@ package examples {
       super.runtimeConfig.copy(shutdownHookTimeout = Duration.Zero)
 
     val run: IO[Unit] =
-      IO(System.exit(0)).uncancelable
+      IO.blocking(System.exit(0)).uncancelable
   }
 
   object FatalErrorUnsafeRun extends IOApp {

tests/jvm/src/test/scala/cats/effect/IOPlatformSpecification.scala

@@ -437,6 +437,7 @@ trait IOPlatformSpecification { self: BaseSpec with ScalaCheck =>
             runtimeBlockingExpiration = 3.seconds,
             reportFailure0 = _.printStackTrace(),
             blockedThreadDetectionEnabled = false,
+            shutdownTimeout = 1.second,
             system = SleepSystem
           )