Fix #4382: Fix thread leak in WSTP by replacing LinkedTransferQueue with SynchronousQueue and ConcurrentHashMap

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

@@ -39,7 +39,7 @@ import scala.concurrent.duration.{Duration, FiniteDuration}
 
 import java.time.Instant
 import java.time.temporal.ChronoField
-import java.util.concurrent.{LinkedTransferQueue, ThreadLocalRandom}
+import java.util.concurrent.{ConcurrentHashMap, SynchronousQueue, ThreadLocalRandom}
 import java.util.concurrent.atomic.{
   AtomicBoolean,
   AtomicInteger,
@@ -131,8 +131,11 @@ private[effect] final class WorkStealingThreadPool[P <: AnyRef](
    */
   private[this] val state: AtomicInteger = new AtomicInteger(threadCount << UnparkShift)
 
-  private[unsafe] val cachedThreads: LinkedTransferQueue[WorkerThread[P]] =
-    new LinkedTransferQueue
+  private[unsafe] val transferStateQueue: SynchronousQueue[WorkerThread.TransferState] =
+    new SynchronousQueue[WorkerThread.TransferState](false)
+
+  private[unsafe] val blockerThreads: ConcurrentHashMap[WorkerThread[P], java.lang.Boolean] =
+    new ConcurrentHashMap()
 
   /**
    * The shutdown latch of the work stealing thread pool.
@@ -749,11 +752,9 @@ private[effect] final class WorkStealingThreadPool[P <: AnyRef](
         system.close()
       }
 
-      var t: WorkerThread[P] = null
-      while ({
-        t = cachedThreads.poll()
-        t ne null
-      }) {
+      val it = blockerThreads.keySet().iterator()
+      while (it.hasNext()) {
+        val t = it.next()
         t.interrupt()
         // don't bother joining, cached threads are not doing anything interesting
       }

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

@@ -26,7 +26,7 @@ import scala.concurrent.{BlockContext, CanAwait}
 import scala.concurrent.duration.{Duration, FiniteDuration}
 
 import java.lang.Long.MIN_VALUE
-import java.util.concurrent.{ArrayBlockingQueue, ThreadLocalRandom}
+import java.util.concurrent.ThreadLocalRandom
 import java.util.concurrent.atomic.AtomicBoolean
 
 import WorkerThread.{Metrics, TransferState}
@@ -110,7 +110,6 @@ private[effect] final class WorkerThread[P <: AnyRef](
    */
   private[this] var _active: Runnable = _
 
-  private val stateTransfer: ArrayBlockingQueue[TransferState] = new ArrayBlockingQueue(1)
   private[this] val runtimeBlockingExpiration: Duration = pool.runtimeBlockingExpiration
 
   private[effect] var currentIOFiber: IOFiber[?] = _
@@ -732,20 +731,27 @@ private[effect] final class WorkerThread[P <: AnyRef](
           // by another thread in the future.
           val len = runtimeBlockingExpiration.length
           val unit = runtimeBlockingExpiration.unit
-          if (pool.cachedThreads.tryTransfer(this, len, unit)) {
-            // Someone accepted the transfer of this thread and will transfer the state soon.
-            val newState = stateTransfer.take()
+
+          // Try to poll for a new state from the transfer queue
+          val newState = pool.transferStateQueue.poll(len, unit)
+
+          if (newState ne null) {
+            // Got a state to take over
             init(newState)
+
           } else {
-            // The timeout elapsed and no one woke up this thread. It's time to exit.
+            // No state to take over after timeout, exit
             pool.blockedWorkerThreadCounter.decrementAndGet()
+            // Remove from blocker threads map if present
+            pool.blockerThreads.remove(this)
             return
           }
         } catch {
           case _: InterruptedException =>
             // This thread was interrupted while cached. This should only happen
             // during the shutdown of the pool. Nothing else to be done, just
             // exit.
+            pool.blockerThreads.remove(this)
             return
         }
       }
@@ -928,15 +934,18 @@ private[effect] final class WorkerThread[P <: AnyRef](
       // Set the name of this thread to a blocker prefixed name.
       setName(s"$prefix-$nameIndex")
 
-      val cached = pool.cachedThreads.poll()
-      if (cached ne null) {
-        // There is a cached worker thread that can be reused.
-        val idx = index
-        pool.replaceWorker(idx, cached)
-        // Transfer the data structures to the cached thread and wake it up.
-        transferState.index = idx
-        transferState.tick = tick + 1
-        val _ = cached.stateTransfer.offer(transferState)
+      val idx = index
+
+      // Prepare the transfer state
+      transferState.index = idx
+      transferState.tick = tick + 1
+
+      // Register this thread in the blockerThreads map
+      val _ = pool.blockerThreads.put(this, java.lang.Boolean.TRUE)
+
+      if (pool.transferStateQueue.offer(transferState)) {
+        // If successful, a waiting thread will pick it up
+
       } else {
         // Spawn a new `WorkerThread`, a literal clone of this one. It is safe to
         // transfer ownership of the local queue and the parked signal to the new
@@ -1002,6 +1011,8 @@ private[effect] final class WorkerThread[P <: AnyRef](
     setName(s"$prefix-${_index}")
 
     blocking = false
+
+    pool.replaceWorker(newIdx, this)
   }
 
   /**