Introduce CoroutineDispatcher.limitedParallelism and make Dispatchers.IO unbounded for limited parallelism (#2918)

* Introduce CoroutineDispatcher.limitedParallelism for granular concurrency control

* Elastic Dispatchers.IO:

    * Extract Ktor-obsolete API to a separate file for backwards compatibility
    * Make Dispatchers.IO being a slice of unlimited blocking scheduler
    * Make Dispatchers.IO.limitParallelism take slices from the same internal scheduler

Fixes #2943
Fixes #2919

Author: qwwdfsad

benchmarks/src/jmh/kotlin/benchmarks/ParametrizedDispatcherBase.kt

@@ -30,7 +30,7 @@ abstract class ParametrizedDispatcherBase : CoroutineScope {
         coroutineContext = when {
             dispatcher == "fjp" -> ForkJoinPool.commonPool().asCoroutineDispatcher()
             dispatcher == "scheduler" -> {
-                ExperimentalCoroutineDispatcher(CORES_COUNT).also { closeable = it }
+                Dispatchers.Default
             }
             dispatcher.startsWith("ftp") -> {
                 newFixedThreadPoolContext(dispatcher.substring(4).toInt(), dispatcher).also { closeable = it }

benchmarks/src/jmh/kotlin/benchmarks/SemaphoreBenchmark.kt

@@ -6,13 +6,10 @@ package benchmarks
 
 import benchmarks.common.*
 import kotlinx.coroutines.*
-import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.scheduling.ExperimentalCoroutineDispatcher
-import kotlinx.coroutines.sync.Semaphore
-import kotlinx.coroutines.sync.withPermit
+import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.sync.*
 import org.openjdk.jmh.annotations.*
-import java.util.concurrent.ForkJoinPool
-import java.util.concurrent.TimeUnit
+import java.util.concurrent.*
 
 @Warmup(iterations = 3, time = 500, timeUnit = TimeUnit.MICROSECONDS)
 @Measurement(iterations = 10, time = 500, timeUnit = TimeUnit.MICROSECONDS)
@@ -84,7 +81,7 @@ open class SemaphoreBenchmark {
 
 enum class SemaphoreBenchDispatcherCreator(val create: (parallelism: Int) -> CoroutineDispatcher) {
     FORK_JOIN({ parallelism -> ForkJoinPool(parallelism).asCoroutineDispatcher() }),
-    EXPERIMENTAL({ parallelism -> ExperimentalCoroutineDispatcher(corePoolSize = parallelism, maxPoolSize = parallelism) })
+    EXPERIMENTAL({ parallelism -> Dispatchers.Default }) // TODO doesn't take parallelism into account
 }
 
 private const val WORK_INSIDE = 80

benchmarks/src/jmh/kotlin/benchmarks/scheduler/actors/PingPongWithBlockingContext.kt

@@ -27,10 +27,8 @@ import kotlin.coroutines.*
 @State(Scope.Benchmark)
 open class PingPongWithBlockingContext {
 
-    @UseExperimental(InternalCoroutinesApi::class)
-    private val experimental = ExperimentalCoroutineDispatcher(8)
-    @UseExperimental(InternalCoroutinesApi::class)
-    private val blocking = experimental.blocking(8)
+    private val experimental = Dispatchers.Default
+    private val blocking = Dispatchers.IO.limitedParallelism(8)
     private val threadPool = newFixedThreadPoolContext(8, "PongCtx")
 
     @TearDown

integration/kotlinx-coroutines-play-services/test/TaskTest.kt

@@ -45,8 +45,8 @@ class TaskTest : TestBase() {
     }
 
     @Test
-    fun testCancelledAsTask() {
-        val deferred = GlobalScope.async {
+    fun testCancelledAsTask() = runTest {
+        val deferred = async(Dispatchers.Default) {
             delay(100)
         }.apply { cancel() }
 
@@ -60,8 +60,8 @@ class TaskTest : TestBase() {
     }
 
     @Test
-    fun testThrowingAsTask() {
-        val deferred = GlobalScope.async<Int> {
+    fun testThrowingAsTask() = runTest({ e -> e is TestException }) {
+        val deferred = async<Int>(Dispatchers.Default) {
             throw TestException("Fail")
         }
 

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -156,6 +156,7 @@ public abstract class kotlinx/coroutines/CoroutineDispatcher : kotlin/coroutines
 	public fun get (Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext$Element;
 	public final fun interceptContinuation (Lkotlin/coroutines/Continuation;)Lkotlin/coroutines/Continuation;
 	public fun isDispatchNeeded (Lkotlin/coroutines/CoroutineContext;)Z
+	public fun limitedParallelism (I)Lkotlinx/coroutines/CoroutineDispatcher;
 	public fun minusKey (Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext;
 	public final fun plus (Lkotlinx/coroutines/CoroutineDispatcher;)Lkotlinx/coroutines/CoroutineDispatcher;
 	public final fun releaseInterceptedContinuation (Lkotlin/coroutines/Continuation;)V
@@ -447,6 +448,7 @@ public class kotlinx/coroutines/JobSupport : kotlinx/coroutines/ChildJob, kotlin
 public abstract class kotlinx/coroutines/MainCoroutineDispatcher : kotlinx/coroutines/CoroutineDispatcher {
 	public fun <init> ()V
 	public abstract fun getImmediate ()Lkotlinx/coroutines/MainCoroutineDispatcher;
+	public fun limitedParallelism (I)Lkotlinx/coroutines/CoroutineDispatcher;
 	public fun toString ()Ljava/lang/String;
 	protected final fun toStringInternalImpl ()Ljava/lang/String;
 }

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

@@ -61,6 +61,45 @@ public abstract class CoroutineDispatcher :
      */
     public open fun isDispatchNeeded(context: CoroutineContext): Boolean = true
 
+    /**
+     * Creates a view of the current dispatcher that limits the parallelism to the given [value][parallelism].
+     * The resulting view uses the original dispatcher for execution, but with the guarantee that
+     * no more than [parallelism] coroutines are executed at the same time.
+     *
+     * This method does not impose restrictions on the number of views or the total sum of parallelism values,
+     * each view controls its own parallelism independently with the guarantee that the effective parallelism
+     * of all views cannot exceed the actual parallelism of the original dispatcher.
+     *
+     * ### Limitations
+     *
+     * The default implementation of `limitedParallelism` does not support direct dispatchers,
+     * such as executing the given runnable in place during [dispatch] calls.
+     * Any dispatcher that may return `false` from [isDispatchNeeded] is considered direct.
+     * For direct dispatchers, it is recommended to override this method
+     * and provide a domain-specific implementation or to throw an [UnsupportedOperationException].
+     *
+     * ### Example of usage
+     * ```
+     * private val backgroundDispatcher = newFixedThreadPoolContext(4, "App Background")
+     * // At most 2 threads will be processing images as it is really slow and CPU-intensive
+     * private val imageProcessingDispatcher = backgroundDispatcher.limitedParallelism(2)
+     * // At most 3 threads will be processing JSON to avoid image processing starvation
+     * private val imageProcessingDispatcher = backgroundDispatcher.limitedParallelism(3)
+     * // At most 1 thread will be doing IO
+     * private val fileWriterDispatcher = backgroundDispatcher.limitedParallelism(1)
+     * ```
+     * is 6. Yet at most 4 coroutines can be executed simultaneously as each view limits only its own parallelism.
+     *
+     * Note that this example was structured in such a way that it illustrates the parallelism guarantees.
+     * In practice, it is usually better to use [Dispatchers.IO] or [Dispatchers.Default] instead of creating a
+     * `backgroundDispatcher`. It is both possible and advised to call `limitedParallelism` on them.
+     */
+    @ExperimentalCoroutinesApi
+    public open fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+        parallelism.checkParallelism()
+        return LimitedDispatcher(this, parallelism)
+    }
+
     /**
      * Dispatches execution of a runnable [block] onto another thread in the given [context].
      * This method should guarantee that the given [block] will be eventually invoked,

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

@@ -115,6 +115,11 @@ internal abstract class EventLoop : CoroutineDispatcher() {
         }
     }
 
+    final override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+        parallelism.checkParallelism()
+        return this
+    }
+
     open fun shutdown() {}
 }
 
@@ -525,4 +530,3 @@ internal expect fun nanoTime(): Long
 internal expect object DefaultExecutor {
     public fun enqueue(task: Runnable)
 }
-

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

@@ -4,6 +4,8 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
+
 /**
  * Base class for special [CoroutineDispatcher] which is confined to application "Main" or "UI" thread
  * and used for any UI-based activities. Instance of `MainDispatcher` can be obtained by [Dispatchers.Main].
@@ -51,6 +53,12 @@ public abstract class MainCoroutineDispatcher : CoroutineDispatcher() {
      */
     override fun toString(): String = toStringInternalImpl() ?: "$classSimpleName@$hexAddress"
 
+    override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+        parallelism.checkParallelism()
+        // MainCoroutineDispatcher is single-threaded -- short-circuit any attempts to limit it
+        return this
+    }
+
     /**
      * Internal method for more specific [toString] implementations. It returns non-null
      * string if this dispatcher is set in the platform as the main one.

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

@@ -11,6 +11,12 @@ import kotlin.jvm.*
  * A coroutine dispatcher that is not confined to any specific thread.
  */
 internal object Unconfined : CoroutineDispatcher() {
+
+    @ExperimentalCoroutinesApi
+    override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+        throw UnsupportedOperationException("limitedParallelism is not supported for Dispatchers.Unconfined")
+    }
+
     override fun isDispatchNeeded(context: CoroutineContext): Boolean = false
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {

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

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+import kotlinx.coroutines.*
+import kotlin.coroutines.*
+import kotlin.jvm.*
+
+/**
+ * The result of .limitedParallelism(x) call, a dispatcher
+ * that wraps the given dispatcher, but limits the parallelism level, while
+ * trying to emulate fairness.
+ */
+internal class LimitedDispatcher(
+    private val dispatcher: CoroutineDispatcher,
+    private val parallelism: Int
+) : CoroutineDispatcher(), Runnable, Delay by (dispatcher as? Delay ?: DefaultDelay) {
+
+    @Volatile
+    private var runningWorkers = 0
+
+    private val queue = LockFreeTaskQueue<Runnable>(singleConsumer = false)
+
+    @ExperimentalCoroutinesApi
+    override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+        parallelism.checkParallelism()
+        if (parallelism >= this.parallelism) return this
+        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
+            }
+
+            @Suppress("CAST_NEVER_SUCCEEDS")
+            synchronized(this as SynchronizedObject) {
+                --runningWorkers
+                if (queue.size == 0) return
+                ++runningWorkers
+                fairnessCounter = 0
+            }
+        }
+    }
+
+    override fun dispatch(context: CoroutineContext, block: Runnable) {
+        dispatchInternal(block) {
+            dispatcher.dispatch(this, this)
+        }
+    }
+
+    @InternalCoroutinesApi
+    override fun dispatchYield(context: CoroutineContext, block: Runnable) {
+        dispatchInternal(block) {
+            dispatcher.dispatchYield(this, this)
+        }
+    }
+
+    private inline fun dispatchInternal(block: Runnable, dispatch: () -> 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()`)
+         */
+        if (!tryAllocateWorker()) return
+        dispatch()
+    }
+
+    private fun tryAllocateWorker(): Boolean {
+        @Suppress("CAST_NEVER_SUCCEEDS")
+        synchronized(this as SynchronizedObject) {
+            if (runningWorkers >= parallelism) return false
+            ++runningWorkers
+            return true
+        }
+    }
+
+    private fun addAndTryDispatching(block: Runnable): Boolean {
+        queue.addLast(block)
+        return runningWorkers >= parallelism
+    }
+}
+
+// Save a few bytecode ops
+internal fun Int.checkParallelism() = require(this >= 1) { "Expected positive parallelism level, but got $this" }