Don't allocate threads on every dispatch in Native's thread pools

kotlinx-coroutines-core/common/src/CloseableCoroutineDispatcher.kt

@@ -19,8 +19,8 @@ public expect abstract class CloseableCoroutineDispatcher() : CoroutineDispatche
 
     /**
      * Initiate the closing sequence of the coroutine dispatcher.
-     * After a successful call to [close], no new tasks will
-     * be accepted to be [dispatched][dispatch], but the previously dispatched tasks will be run.
+     * After a successful call to [close], no new tasks will be accepted to be [dispatched][dispatch].
+     * The previously-submitted tasks will still be run, but [close] is not guaranteed to wait for them to finish.
      *
      * Invocations of `close` are idempotent and thread-safe.
      */

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlinx.atomicfu.*
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class MultithreadedDispatcherStressTest {
+    val shared = atomic(0)
+
+    /**
+     * Tests that [newFixedThreadPoolContext] will not drop tasks when closed.
+     */
+    @Test
+    fun testClosingNotDroppingTasks() {
+        repeat(7) {
+            shared.value = 0
+            val nThreads = it + 1
+            val dispatcher = newFixedThreadPoolContext(nThreads, "testMultiThreadedContext")
+            repeat(1_000) {
+                dispatcher.dispatch(EmptyCoroutineContext, Runnable {
+                    shared.incrementAndGet()
+                })
+            }
+            dispatcher.close()
+            while (shared.value < 1_000) {
+                // spin.
+                // the test will hang here if the dispatcher drops tasks.
+            }
+        }
+    }
+}

kotlinx-coroutines-core/native/src/MultithreadedDispatchers.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
@@ -73,43 +74,91 @@ private class MultiWorkerDispatcher(
     workersCount: Int
 ) : CloseableCoroutineDispatcher() {
     private val tasksQueue = Channel<Runnable>(Channel.UNLIMITED)
+    private val availableWorkers = Channel<CancellableContinuation<Runnable>>(Channel.UNLIMITED)
     private val workerPool = OnDemandAllocatingPool(workersCount) {
         Worker.start(name = "$name-$it").apply {
             executeAfter { workerRunLoop() }
         }
     }
 
+    /**
+     * (number of tasks - number of workers) * 2 + (1 if closed)
+     */
+    private val tasksAndWorkersCounter = atomic(0L)
+
+    private inline fun Long.isClosed() = this and 1L == 1L
+    private inline fun Long.hasTasks() = this >= 2
+    private inline fun Long.hasWorkers() = this < 0
+
     private fun workerRunLoop() = runBlocking {
-        // NB: we leverage tail-call optimization in this loop, do not replace it with
-        // .receive() without proper evaluation
-        for (task in tasksQueue) {
-            /**
-             * Any unhandled exception here will pass through worker's boundary and will be properly reported.
-             */
-            task.run()
+        while (true) {
+            val state = tasksAndWorkersCounter.getAndUpdate {
+                if (it.isClosed() && !it.hasTasks()) return@runBlocking
+                it - 2
+            }
+            if (state.hasTasks()) {
+                // we promised to process a task, and there are some
+                tasksQueue.receive().run()
+            } else {
+                try {
+                    suspendCancellableCoroutine {
+                        val result = availableWorkers.trySend(it)
+                        checkChannelResult(result)
+                    }.run()
+                } catch (e: CancellationException) {
+                    /** we are cancelled from [close] and thus will never get back to this branch of code,
+                    but there may still be pending work, so we can't just exit here. */
+                }
+            }
         }
     }
 
+    // a worker that promised to be here and should actually arrive, so we wait for it in a blocking manner.
+    private fun obtainWorker(): CancellableContinuation<Runnable> =
+        availableWorkers.tryReceive().getOrNull() ?: runBlocking { availableWorkers.receive() }
+
     override fun dispatch(context: CoroutineContext, block: Runnable) {
-        fun throwClosed(block: Runnable) {
-            throw IllegalStateException("Dispatcher $name was closed, attempted to schedule: $block")
+        val state = tasksAndWorkersCounter.getAndUpdate {
+            if (it.isClosed())
+                throw IllegalStateException("Dispatcher $name was closed, attempted to schedule: $block")
+            it + 2
         }
-
-        if (!workerPool.allocate()) throwClosed(block) // Do not even try to send to avoid race
-
-        tasksQueue.trySend(block).onClosed {
-            throwClosed(block)
+        if (state.hasWorkers()) {
+            // there are workers that have nothing to do, let's grab one of them
+            obtainWorker().resume(block)
+        } else {
+            workerPool.allocate()
+            // no workers are available, we must queue the task
+            val result = tasksQueue.trySend(block)
+            checkChannelResult(result)
         }
     }
 
     override fun close() {
-        val workers = workerPool.close()
-        tasksQueue.close()
+        tasksAndWorkersCounter.getAndUpdate { if (it.isClosed()) it else it or 1L }
+        val workers = workerPool.close() // no new workers will be created
+        while (true) {
+            // check if there are workers that await tasks in their personal channels, we need to wake them up
+            val state = tasksAndWorkersCounter.getAndUpdate {
+                if (it.hasWorkers()) it + 2 else it
+            }
+            if (!state.hasWorkers())
+                break
+            obtainWorker().cancel()
+        }
         /*
          * Here we cannot avoid waiting on `.result`, otherwise it will lead
          * to a native memory leak, including a pthread handle.
          */
         val requests = workers.map { it.requestTermination() }
         requests.map { it.result }
     }
+
+    private fun checkChannelResult(result: ChannelResult<*>) {
+        if (!result.isSuccess)
+            throw IllegalStateException(
+                "Internal invariants of $this were violated, please file a bug to kotlinx.coroutines",
+                result.exceptionOrNull()
+            )
+    }
 }

kotlinx-coroutines-core/native/test/MultithreadedDispatchersTest.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlin.native.concurrent.*
+import kotlin.test.*
+
+class MultithreadedDispatchersTest {
+    /**
+     * Test that [newFixedThreadPoolContext] does not allocate more dispatchers than it needs to.
+     * Incidentally also tests that it will allocate enough workers for its needs. Otherwise, the test will hang.
+     */
+    @Test
+    fun testNotAllocatingExtraDispatchers() {
+        suspend fun spin(set: MutableSet<Worker>) {
+            repeat(100) {
+                set.add(Worker.current)
+                delay(1)
+            }
+        }
+        val dispatcher = newFixedThreadPoolContext(64, "test")
+        try {
+            runBlocking {
+                val encounteredWorkers = mutableSetOf<Worker>()
+                var canStart1 = false
+                var canStart2 = false
+                val coroutine1 = launch(dispatcher) {
+                    while (!canStart1) {
+                        // intentionally empty
+                    }
+                    canStart2 = true
+                    spin(encounteredWorkers)
+                }
+                val coroutine2 = launch(dispatcher) {
+                    canStart1 = true
+                    while (!canStart2) {
+                        // intentionally empty
+                    }
+                    spin(encounteredWorkers)
+                }
+                listOf(coroutine1, coroutine2).joinAll()
+                assertEquals(2, encounteredWorkers.size)
+            }
+        } finally {
+            dispatcher.close()
+        }
+    }
+}