Fix oversubscription in Dispatchers.Default (#3684)

Previously, an arbitrary task was added to the scheduler's queue and the corresponding counter was incremented. The actual executing thread decremented the counter as soon as the task was done, meaning that, if timings are unlucky enough, the overall balance of blocking tasks could've been negative.

The second part of the equation is that all increments are atomic (instead of being CAS-based), meaning that due to programmatic bugs like the previous one, it's easy to overstep the boundary of adjacent masks, leaving the state in an inconsistent state.

Fix is trivial -- increment counter before publishing, do not reset it if the scheduler is closed and task wasn't added.

Fixes #3642

Author: qwwdfsad

kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt

@@ -265,9 +265,13 @@ internal class CoroutineScheduler(
 
     /**
      * The `Long` value describing the state of workers in this pool.
-     * Currently includes created, CPU-acquired, and blocking workers, each occupying [BLOCKING_SHIFT] bits.
+     * Currently, includes created, CPU-acquired, and blocking workers, each occupying [BLOCKING_SHIFT] bits.
+     *
+     * State layout (highest to lowest bits):
+     * | --- number of cpu permits, 22 bits ---  | --- number of blocking tasks, 21 bits ---  | --- number of created threads, 21 bits ---  |
      */
     private val controlState = atomic(corePoolSize.toLong() shl CPU_PERMITS_SHIFT)
+
     private val createdWorkers: Int inline get() = (controlState.value and CREATED_MASK).toInt()
     private val availableCpuPermits: Int inline get() = availableCpuPermits(controlState.value)
 
@@ -383,6 +387,10 @@ internal class CoroutineScheduler(
     fun dispatch(block: Runnable, taskContext: TaskContext = NonBlockingContext, tailDispatch: Boolean = false) {
         trackTask() // this is needed for virtual time support
         val task = createTask(block, taskContext)
+        val isBlockingTask = task.isBlocking
+        // Invariant: we increment counter **before** publishing the task
+        // so executing thread can safely decrement the number of blocking tasks
+        val stateSnapshot = if (isBlockingTask) incrementBlockingTasks() else 0
         // try to submit the task to the local queue and act depending on the result
         val currentWorker = currentWorker()
         val notAdded = currentWorker.submitToLocalQueue(task, tailDispatch)
@@ -394,12 +402,12 @@ internal class CoroutineScheduler(
         }
         val skipUnpark = tailDispatch && currentWorker != null
         // Checking 'task' instead of 'notAdded' is completely okay
-        if (task.mode == TASK_NON_BLOCKING) {
+        if (isBlockingTask) {
+            // Use state snapshot to better estimate the number of running threads
+            signalBlockingWork(stateSnapshot, skipUnpark = skipUnpark)
+        } else {
             if (skipUnpark) return
             signalCpuWork()
-        } else {
-            // Increment blocking tasks anyway
-            signalBlockingWork(skipUnpark = skipUnpark)
         }
     }
 
@@ -413,11 +421,11 @@ internal class CoroutineScheduler(
         return TaskImpl(block, nanoTime, taskContext)
     }
 
-    private fun signalBlockingWork(skipUnpark: Boolean) {
-        // Use state snapshot to avoid thread overprovision
-        val stateSnapshot = incrementBlockingTasks()
+    // NB: should only be called from 'dispatch' method due to blocking tasks increment
+    private fun signalBlockingWork(stateSnapshot: Long, skipUnpark: Boolean) {
         if (skipUnpark) return
         if (tryUnpark()) return
+        // Use state snapshot to avoid accidental thread overprovision
         if (tryCreateWorker(stateSnapshot)) return
         tryUnpark() // Try unpark again in case there was race between permit release and parking
     }

kotlinx-coroutines-core/jvm/test/scheduling/CoroutineSchedulerInternalApiStressTest.kt

@@ -80,8 +80,8 @@ class CoroutineSchedulerInternalApiStressTest : TestBase() {
                     }
                 }
                 completionLatch.countDown()
-//                assertEquals(100, timesHelped)
-//                assertTrue(Thread.currentThread() in observedDefaultThreads, observedDefaultThreads.toString())
+                assertEquals(100, timesHelped)
+                assertTrue(Thread.currentThread() in observedDefaultThreads, observedDefaultThreads.toString())
             }
         }
     }