Fix oversubscription in CoroutineScheduler (-> Dispatchers.Default)

Previously, a worker thread unconditionally processed tasks from its own local queue, even if tasks were CPU-intensive, but CPU token was not acquired.

Fixes #3418

Author: qwwdfsad

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

@@ -874,8 +874,14 @@ internal class CoroutineScheduler(
 
         fun findTask(mayHaveLocalTasks: Boolean): Task? {
             if (tryAcquireCpuPermit()) return findAnyTask(mayHaveLocalTasks)
-            // If we can't acquire a CPU permit -- attempt to find blocking task
-            return globalBlockingQueue.removeFirstOrNull() ?: trySteal(blockingOnly = true)
+            /*
+             * If we can't acquire a CPU permit, attempt to find blocking task:
+             * * Check if our queue has one (maybe mixed in with CPU tasks)
+             * * Poll global and try steal
+             */
+            return localQueue.pollBlocking()
+                ?: globalBlockingQueue.removeFirstOrNull()
+                ?: trySteal(blockingOnly = true)
         }
 
         private fun findAnyTask(scanLocalQueue: Boolean): Task? {

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

@@ -47,10 +47,12 @@ internal class WorkQueue {
      * [T2] changeProducerIndex (3)
      * [T3] changeConsumerIndex (4)
      *
-     * Which can lead to resulting size bigger than actual size at any moment of time.
-     * This is in general harmless because steal will be blocked by timer
+     * Which can lead to resulting size being negative or bigger than actual size at any moment of time.
+     * This is in general harmless because steal will be blocked by timer.
+     * Negative sizes can be observed only when non-owner reads the size, which happens only
+     * for diagnostic toString().
      */
-    internal val bufferSize: Int get() = producerIndex.value - consumerIndex.value
+    private val bufferSize: Int get() = producerIndex.value - consumerIndex.value
     internal val size: Int get() = if (lastScheduledTask.value != null) bufferSize + 1 else bufferSize
     private val buffer: AtomicReferenceArray<Task?> = AtomicReferenceArray(BUFFER_CAPACITY)
     private val lastScheduledTask = atomic<Task?>(null)
@@ -134,6 +136,32 @@ internal class WorkQueue {
         return tryStealLastScheduled(stolenTaskRef, blockingOnly = true)
     }
 
+    // Polls for blocking task, invoked only by the owner
+    fun pollBlocking(): Task? {
+        while (true) { // Poll the slot
+            val lastScheduled = lastScheduledTask.value ?: break
+            if (!lastScheduled.isBlocking) break
+            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
+                return lastScheduled
+            } // Failed -> someone else stole it
+        }
+
+        val start = consumerIndex.value
+        var end = producerIndex.value
+
+        while (start != end) {
+            --end
+            val index = end and MASK
+            if (blockingTasksInBuffer.value == 0) break
+            val value = buffer[index]
+            if (value != null && value.isBlocking && buffer.compareAndSet(index, value, null)) {
+                blockingTasksInBuffer.decrementAndGet()
+                return value
+            }
+        }
+        return null
+    }
+
     fun offloadAllWorkTo(globalQueue: GlobalQueue) {
         lastScheduledTask.getAndSet(null)?.let { globalQueue.addLast(it) }
         while (pollTo(globalQueue)) {

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

@@ -41,7 +41,7 @@ class WorkQueueStressTest : TestBase() {
         threads += thread(name = "producer") {
             startLatch.await()
             for (i in 1..offerIterations) {
-                while (producerQueue.bufferSize > BUFFER_CAPACITY / 2) {
+                while (producerQueue.size > BUFFER_CAPACITY / 2) {
                     Thread.yield()
                 }
 
@@ -79,7 +79,7 @@ class WorkQueueStressTest : TestBase() {
         threads += thread(name = "producer") {
             startLatch.await()
             for (i in 1..offerIterations) {
-                while (producerQueue.bufferSize == BUFFER_CAPACITY - 1) {
+                while (producerQueue.size == BUFFER_CAPACITY - 1) {
                     Thread.yield()
                 }
 

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

@@ -69,9 +69,20 @@ class WorkQueueTest : TestBase() {
         assertEquals(TASK_STOLEN, victim.trySteal(ref))
         assertEquals(arrayListOf(2L), stealer.drain(ref))
     }
+
+    @Test
+    fun testPollBlocking() {
+        val queue = WorkQueue()
+        assertNull(queue.pollBlocking())
+        val blockingTask = blockingTask(1L)
+        queue.add(blockingTask)
+        queue.add(task(1L))
+        assertSame(blockingTask, queue.pollBlocking())
+    }
 }
 
 internal fun task(n: Long) = TaskImpl(Runnable {}, n, NonBlockingContext)
+internal fun blockingTask(n: Long) = TaskImpl(Runnable {}, n, BlockingContext)
 
 internal fun WorkQueue.drain(ref: ObjectRef<Task?>): List<Long> {
     var task: Task? = poll()