New implementation of K/N multi-threaded dispatchers (#3421)

* Get rid of the race on LimitedDispatcher due to @volatile being no-op on K/N

Improve newFixedThreadPoolContext on K/N:

* Create workers lazily
* Do properly handle rejection
* Await termination in parallel
* Extract the concurrent structure from MultithreadedDispatchers to test it with Lincheck

* Fix a race

Co-authored-by: dkhalanskyjb <[email protected]>
Co-authored-by: Dmitry Khalanskiy <[email protected]>

Author: qwwdfsad

kotlinx-coroutines-core/build.gradle

@@ -185,7 +185,7 @@ jvmTest {
     }
     // 'stress' is required to be able to run all subpackage tests like ":jvmTests --tests "*channels*" -Pstress=true"
     if (!Idea.active && rootProject.properties['stress'] == null) {
-        exclude '**/*LincheckTest.*'
+        exclude '**/*LincheckTest*'
         exclude '**/*StressTest.*'
     }
     if (Idea.active) {
@@ -229,7 +229,7 @@ task jvmStressTest(type: Test, dependsOn: compileTestKotlinJvm) {
 task jvmLincheckTest(type: Test, dependsOn: compileTestKotlinJvm) {
     classpath = files { jvmTest.classpath }
     testClassesDirs = files { jvmTest.testClassesDirs }
-    include '**/*LincheckTest.*'
+    include '**/*LincheckTest*'
     enableAssertions = true
     testLogging.showStandardStreams = true
     configureJvmForLincheck(jvmLincheckTest)

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

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.internal
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
@@ -18,8 +19,9 @@ internal class LimitedDispatcher(
     private val parallelism: Int
 ) : CoroutineDispatcher(), Runnable, Delay by (dispatcher as? Delay ?: DefaultDelay) {
 
-    @Volatile
-    private var runningWorkers = 0
+    // Atomic is necessary here for the sake of K/N memory ordering,
+    // there is no need in atomic operations for this property
+    private val runningWorkers = atomic(0)
 
     private val queue = LockFreeTaskQueue<Runnable>(singleConsumer = false)
 
@@ -54,9 +56,9 @@ internal class LimitedDispatcher(
             }
 
             synchronized(workerAllocationLock) {
-                --runningWorkers
+                runningWorkers.decrementAndGet()
                 if (queue.size == 0) return
-                ++runningWorkers
+                runningWorkers.incrementAndGet()
                 fairnessCounter = 0
             }
         }
@@ -90,15 +92,15 @@ internal class LimitedDispatcher(
 
     private fun tryAllocateWorker(): Boolean {
         synchronized(workerAllocationLock) {
-            if (runningWorkers >= parallelism) return false
-            ++runningWorkers
+            if (runningWorkers.value >= parallelism) return false
+            runningWorkers.incrementAndGet()
             return true
         }
     }
 
     private fun addAndTryDispatching(block: Runnable): Boolean {
         queue.addLast(block)
-        return runningWorkers >= parallelism
+        return runningWorkers.value >= parallelism
     }
 }
 

kotlinx-coroutines-core/concurrent/src/internal/OnDemandAllocatingPool.kt

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2016-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+import kotlinx.atomicfu.*
+
+/**
+ * A thread-safe resource pool.
+ *
+ * [maxCapacity] is the maximum amount of elements.
+ * [create] is the function that creates a new element.
+ *
+ * This is only used in the Native implementation,
+ * but is part of the `concurrent` source set in order to test it on the JVM.
+ */
+internal class OnDemandAllocatingPool<T>(
+    private val maxCapacity: Int,
+    private val create: (Int) -> T
+) {
+    /**
+     * Number of existing elements + isClosed flag in the highest bit.
+     * Once the flag is set, the value is guaranteed not to change anymore.
+     */
+    private val controlState = atomic(0)
+    private val elements = atomicArrayOfNulls<T>(maxCapacity)
+
+    /**
+     * Returns the number of elements that need to be cleaned up due to the pool being closed.
+     */
+    @Suppress("NOTHING_TO_INLINE")
+    private inline fun tryForbidNewElements(): Int {
+        controlState.loop {
+            if (it.isClosed()) return 0 // already closed
+            if (controlState.compareAndSet(it, it or IS_CLOSED_MASK)) return it
+        }
+    }
+
+    @Suppress("NOTHING_TO_INLINE")
+    private inline fun Int.isClosed(): Boolean = this and IS_CLOSED_MASK != 0
+
+    /**
+     * Request that a new element is created.
+     *
+     * Returns `false` if the pool is closed.
+     *
+     * Note that it will still return `true` even if an element was not created due to reaching [maxCapacity].
+     *
+     * Rethrows the exceptions thrown from [create]. In this case, this operation has no effect.
+     */
+    fun allocate(): Boolean {
+        controlState.loop { ctl ->
+            if (ctl.isClosed()) return false
+            if (ctl >= maxCapacity) return true
+            if (controlState.compareAndSet(ctl, ctl + 1)) {
+                elements[ctl].value = create(ctl)
+                return true
+            }
+        }
+    }
+
+    /**
+     * Close the pool.
+     *
+     * This will prevent any new elements from being created.
+     * All the elements present in the pool will be returned.
+     *
+     * The function is thread-safe.
+     *
+     * [close] can be called multiple times, but only a single call will return a non-empty list.
+     * This is due to the elements being cleaned out from the pool on the first invocation to avoid memory leaks,
+     * and no new elements being created after.
+     */
+    fun close(): List<T> {
+        val elementsExisting = tryForbidNewElements()
+        return (0 until elementsExisting).map { i ->
+            // we wait for the element to be created, because we know that eventually it is going to be there
+            loop {
+                val element = elements[i].getAndSet(null)
+                if (element != null) {
+                    return@map element
+                }
+            }
+        }
+    }
+
+    // for tests
+    internal fun stateRepresentation(): String {
+        val ctl = controlState.value
+        val elementsStr = (0 until (ctl and IS_CLOSED_MASK.inv())).map { elements[it].value }.toString()
+        val closedStr = if (ctl.isClosed()) "[closed]" else ""
+        return elementsStr + closedStr
+    }
+
+    override fun toString(): String = "OnDemandAllocatingPool(${stateRepresentation()})"
+}
+
+// KT-25023
+private inline fun loop(block: () -> Unit): Nothing {
+    while (true) {
+        block()
+    }
+}
+
+private const val IS_CLOSED_MASK = 1 shl 31

kotlinx-coroutines-core/jvm/test/internal/OnDemandAllocatingPoolLincheckTest.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2016-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+import kotlinx.atomicfu.*
+import kotlinx.coroutines.*
+import org.jetbrains.kotlinx.lincheck.*
+import org.jetbrains.kotlinx.lincheck.annotations.*
+
+/**
+ * Test that:
+ * * All elements allocated in [OnDemandAllocatingPool] get returned when [close] is invoked.
+ * * After reaching the maximum capacity, new elements are not added.
+ * * After [close] is invoked, [OnDemandAllocatingPool.allocate] returns `false`.
+ * * [OnDemandAllocatingPool.close] will return an empty list after the first invocation.
+ */
+abstract class OnDemandAllocatingPoolLincheckTest(maxCapacity: Int) : AbstractLincheckTest() {
+    private val counter = atomic(0)
+    private val pool = OnDemandAllocatingPool(maxCapacity = maxCapacity, create = {
+        counter.getAndIncrement()
+    })
+
+    @Operation
+    fun allocate(): Boolean = pool.allocate()
+
+    @Operation
+    fun close(): String = pool.close().sorted().toString()
+
+    override fun extractState(): Any = pool.stateRepresentation()
+}
+
+abstract class OnDemandAllocatingSequentialPool(private val maxCapacity: Int) {
+    var closed = false
+    var elements = 0
+
+    fun allocate() = if (closed) {
+        false
+    } else {
+        if (elements < maxCapacity) {
+            elements++
+        }
+        true
+    }
+
+    fun close(): String = if (closed) {
+        emptyList()
+    } else {
+        closed = true
+        (0 until elements)
+    }.sorted().toString()
+}
+
+class OnDemandAllocatingPool3LincheckTest : OnDemandAllocatingPoolLincheckTest(3) {
+    override fun <O : Options<O, *>> O.customize(isStressTest: Boolean): O =
+        this.sequentialSpecification(OnDemandAllocatingSequentialPool3::class.java)
+}
+
+class OnDemandAllocatingSequentialPool3 : OnDemandAllocatingSequentialPool(3)

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

@@ -5,6 +5,7 @@
 package kotlinx.coroutines
 
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 import kotlin.native.concurrent.*
 
@@ -14,7 +15,7 @@ public actual fun newSingleThreadContext(name: String): CloseableCoroutineDispat
 }
 
 public actual fun newFixedThreadPoolContext(nThreads: Int, name: String): CloseableCoroutineDispatcher {
-    require(nThreads >= 1) { "Expected at least one thread, but got: $nThreads"}
+    require(nThreads >= 1) { "Expected at least one thread, but got: $nThreads" }
     return MultiWorkerDispatcher(name, nThreads)
 }
 
@@ -67,28 +68,48 @@ internal class WorkerDispatcher(name: String) : CloseableCoroutineDispatcher(),
     }
 }
 
-private class MultiWorkerDispatcher(name: String, workersCount: Int) : CloseableCoroutineDispatcher() {
+private class MultiWorkerDispatcher(
+    private val name: String,
+    workersCount: Int
+) : CloseableCoroutineDispatcher() {
     private val tasksQueue = Channel<Runnable>(Channel.UNLIMITED)
-    private val workers = Array(workersCount) { Worker.start(name = "$name-$it") }
-
-    init {
-        workers.forEach { w -> w.executeAfter(0L) { workerRunLoop() } }
+    private val workerPool = OnDemandAllocatingPool(workersCount) {
+        Worker.start(name = "$name-$it").apply {
+            executeAfter { workerRunLoop() }
+        }
     }
 
     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) {
-            // TODO error handling
+            /**
+             * Any unhandled exception here will pass through worker's boundary and will be properly reported.
+             */
             task.run()
         }
     }
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
-        // TODO handle rejections
-        tasksQueue.trySend(block)
+        fun throwClosed(block: Runnable) {
+            throw IllegalStateException("Dispatcher $name was closed, attempted to schedule: $block")
+        }
+
+        if (!workerPool.allocate()) throwClosed(block) // Do not even try to send to avoid race
+
+        tasksQueue.trySend(block).onClosed {
+            throwClosed(block)
+        }
     }
 
     override fun close() {
+        val workers = workerPool.close()
         tasksQueue.close()
-        workers.forEach { it.requestTermination().result }
+        /*
+         * 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 }
     }
 }