Fix limitedParallelism doing dispatches when it has no tasks (#3672)

To support fairness w.r.t. the tasks that want to run on the
wrapped dispatcher outside of its view, `LimitedDispatcher`
redispatches the work loop periodically. This can happen even when
there are no more tasks in `LimitedDispatcher`'s queue.
Typically, this behavior is completely benign: after being
redistpatched, the new runner just checks the queue, sees that
there is nothing there, and leaves. However, with closable
dispatchers, this affects correctness, as this redispatch may
happen even after all the dispatched tasks were run and the
dispatcher was closed.

Author: dkhalanskyjb

kotlinx-coroutines-core/common/src/internal/LimitedDispatcher.kt

@@ -7,7 +7,6 @@ package kotlinx.coroutines.internal
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
-import kotlin.jvm.*
 
 /**
  * The result of .limitedParallelism(x) call, a dispatcher
@@ -27,7 +26,7 @@ import kotlin.jvm.*
 internal class LimitedDispatcher(
     private val dispatcher: CoroutineDispatcher,
     private val parallelism: Int
-) : CoroutineDispatcher(), Runnable, Delay by (dispatcher as? Delay ?: DefaultDelay) {
+) : CoroutineDispatcher(), Delay by (dispatcher as? Delay ?: DefaultDelay) {
 
     // Atomic is necessary here for the sake of K/N memory ordering,
     // there is no need in atomic operations for this property
@@ -45,61 +44,37 @@ internal class LimitedDispatcher(
         return super.limitedParallelism(parallelism)
     }
 
-    override fun run() {
-        var fairnessCounter = 0
-        while (true) {
-            val task = queue.removeFirstOrNull()
-            if (task != null) {
-                try {
-                    task.run()
-                } catch (e: Throwable) {
-                    handleCoroutineException(EmptyCoroutineContext, e)
-                }
-                // 16 is our out-of-thin-air constant to emulate fairness. Used in JS dispatchers as well
-                if (++fairnessCounter >= 16 && dispatcher.isDispatchNeeded(this)) {
-                    // Do "yield" to let other views to execute their runnable as well
-                    // Note that we do not decrement 'runningWorkers' as we still committed to do our part of work
-                    dispatcher.dispatch(this, this)
-                    return
-                }
-                continue
-            }
-
-            synchronized(workerAllocationLock) {
-                runningWorkers.decrementAndGet()
-                if (queue.size == 0) return
-                runningWorkers.incrementAndGet()
-                fairnessCounter = 0
-            }
-        }
-    }
-
     override fun dispatch(context: CoroutineContext, block: Runnable) {
-        dispatchInternal(block) {
-            dispatcher.dispatch(this, this)
+        dispatchInternal(block) { worker ->
+            dispatcher.dispatch(this, worker)
         }
     }
 
     @InternalCoroutinesApi
     override fun dispatchYield(context: CoroutineContext, block: Runnable) {
-        dispatchInternal(block) {
-            dispatcher.dispatchYield(this, this)
+        dispatchInternal(block) { worker ->
+            dispatcher.dispatchYield(this, worker)
         }
     }
 
-    private inline fun dispatchInternal(block: Runnable, dispatch: () -> Unit) {
+    /**
+     * Tries to dispatch the given [block].
+     * If there are not enough workers, it starts a new one via [startWorker].
+     */
+    private inline fun dispatchInternal(block: Runnable, startWorker: (Worker) -> Unit) {
         // Add task to queue so running workers will be able to see that
-        if (addAndTryDispatching(block)) return
-        /*
-         * Protect against the race when the number of workers is enough,
-         * but one (because of synchronized serialization) attempts to complete,
-         * and we just observed the number of running workers smaller than the actual
-         * number (hit right between `--runningWorkers` and `++runningWorkers` in `run()`)
-         */
+        queue.addLast(block)
+        if (runningWorkers.value >= parallelism) return
+        // allocation may fail if some workers were launched in parallel or a worker temporarily decreased
+        // `runningWorkers` when they observed an empty queue.
         if (!tryAllocateWorker()) return
-        dispatch()
+        val task = obtainTaskOrDeallocateWorker() ?: return
+        startWorker(Worker(task))
     }
 
+    /**
+     * Tries to obtain the permit to start a new worker.
+     */
     private fun tryAllocateWorker(): Boolean {
         synchronized(workerAllocationLock) {
             if (runningWorkers.value >= parallelism) return false
@@ -108,9 +83,49 @@ internal class LimitedDispatcher(
         }
     }
 
-    private fun addAndTryDispatching(block: Runnable): Boolean {
-        queue.addLast(block)
-        return runningWorkers.value >= parallelism
+    /**
+     * Obtains the next task from the queue, or logically deallocates the worker if the queue is empty.
+     */
+    private fun obtainTaskOrDeallocateWorker(): Runnable? {
+        while (true) {
+            when (val nextTask = queue.removeFirstOrNull()) {
+                null -> synchronized(workerAllocationLock) {
+                    runningWorkers.decrementAndGet()
+                    if (queue.size == 0) return null
+                    runningWorkers.incrementAndGet()
+                }
+                else -> return nextTask
+            }
+        }
+    }
+
+    /**
+     * A worker that polls the queue and runs tasks until there are no more of them.
+     *
+     * It always stores the next task to run. This is done in order to prevent the possibility of the fairness
+     * re-dispatch happening when there are no more tasks in the queue. This is important because, after all the
+     * actual tasks are done, nothing prevents the user from closing the dispatcher and making it incorrect to
+     * perform any more dispatches.
+     */
+    private inner class Worker(private var currentTask: Runnable) : Runnable {
+        override fun run() {
+            var fairnessCounter = 0
+            while (true) {
+                try {
+                    currentTask.run()
+                } catch (e: Throwable) {
+                    handleCoroutineException(EmptyCoroutineContext, e)
+                }
+                currentTask = obtainTaskOrDeallocateWorker() ?: return
+                // 16 is our out-of-thin-air constant to emulate fairness. Used in JS dispatchers as well
+                if (++fairnessCounter >= 16 && dispatcher.isDispatchNeeded(this@LimitedDispatcher)) {
+                    // Do "yield" to let other views execute their runnable as well
+                    // Note that we do not decrement 'runningWorkers' as we are still committed to our part of work
+                    dispatcher.dispatch(this@LimitedDispatcher, this)
+                    return
+                }
+            }
+        }
     }
 }
 

kotlinx-coroutines-core/concurrent/test/LimitedParallelismConcurrentTest.kt

@@ -7,6 +7,7 @@ package kotlinx.coroutines
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.exceptions.*
+import kotlin.coroutines.*
 import kotlin.test.*
 
 class LimitedParallelismConcurrentTest : TestBase() {
@@ -58,4 +59,37 @@ class LimitedParallelismConcurrentTest : TestBase() {
         joinAll(j1, j2)
         executor.close()
     }
+
+    /**
+     * Tests that, when no tasks are present, the limited dispatcher does not dispatch any tasks.
+     * This is important for the case when a dispatcher is closeable and the [CoroutineDispatcher.limitedParallelism]
+     * machinery could trigger a dispatch after the dispatcher is closed.
+     */
+    @Test
+    fun testNotDoingDispatchesWhenNoTasksArePresent() = runTest {
+        class NaggingDispatcher: CoroutineDispatcher() {
+            val closed = atomic(false)
+            override fun dispatch(context: CoroutineContext, block: Runnable) {
+                if (closed.value)
+                    fail("Dispatcher was closed, but still dispatched a task")
+                Dispatchers.Default.dispatch(context, block)
+            }
+            fun close() {
+                closed.value = true
+            }
+        }
+        repeat(stressTestMultiplier * 500_000) {
+            val dispatcher = NaggingDispatcher()
+            val view = dispatcher.limitedParallelism(1)
+            val deferred = CompletableDeferred<Unit>()
+            val job = launch(view) {
+                deferred.await()
+            }
+            launch(Dispatchers.Default) {
+                deferred.complete(Unit)
+            }
+            job.join()
+            dispatcher.close()
+        }
+    }
 }

kotlinx-coroutines-core/jvm/test/TestBase.kt

@@ -97,10 +97,10 @@ public actual open class TestBase(private var disableOutCheck: Boolean)  {
 
     private fun printError(message: String, cause: Throwable) {
         setError(cause)
-        println("$message: $cause")
-        cause.printStackTrace(System.out)
-        println("--- Detected at ---")
-        Throwable().printStackTrace(System.out)
+        System.err.println("$message: $cause")
+        cause.printStackTrace(System.err)
+        System.err.println("--- Detected at ---")
+        Throwable().printStackTrace(System.err)
     }
 
     /**