Preliminary prototype of #3439

qwwdfsad · qwwdfsad · commit 8cee1fdb2d80 · 2023-01-03T14:54:24.000+01:00
diff --git a/kotlinx-coroutines-core/common/src/internal/LimitedDispatcher.kt b/kotlinx-coroutines-core/common/src/internal/LimitedDispatcher.kt
@@ -13,6 +13,16 @@ 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.
+ *
+ * ### Implementation details
+ *
+ * By design, 'LimitedDispatcher' never [dispatches][CoroutineDispatcher.dispatch] originally sent tasks
+ * to the underlying dispatcher. Instead, it maintains its own queue of tasks sent to this dispatcher and
+ * dispatches at most [parallelism] "worker-loop" tasks that poll the underlying queue and cooperatively preempt
+ * in order to avoid starvation of the underlying dispatcher.
+ *
+ * Such invariant is crucial in order to be compatible with any underlying dispatcher implementation without
+ * direct cooperation.
  */
 internal class LimitedDispatcher(
     private val dispatcher: CoroutineDispatcher,
diff --git a/kotlinx-coroutines-core/jvm/src/EventLoop.kt b/kotlinx-coroutines-core/jvm/src/EventLoop.kt
@@ -4,6 +4,10 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.Runnable
+import kotlinx.coroutines.scheduling.*
+import kotlinx.coroutines.scheduling.CoroutineScheduler
+
 internal actual abstract class EventLoopImplPlatform: EventLoop() {
     protected abstract val thread: Thread
 
@@ -45,6 +49,62 @@ internal actual fun createEventLoop(): EventLoop = BlockingEventLoop(Thread.curr
  */
 @InternalCoroutinesApi
 public fun processNextEventInCurrentThread(): Long =
+    // This API is used in Ktor for serverless integration where a single thread awaits a blocking call
+    // (and, to avoid actual blocking, does something via this call), see #850
     ThreadLocalEventLoop.currentOrNull()?.processNextEvent() ?: Long.MAX_VALUE
 
 internal actual inline fun platformAutoreleasePool(crossinline block: () -> Unit) = block()
+
+/**
+ * Retrieves and executes a single task from the current system dispatcher ([Dispatchers.Default] or [Dispatchers.IO]).
+ * Returns `0` if any task was executed, `>= 0` for number of nanoseconds to wait until invoking this method again
+ * (implying that there will be a task to steal in N nanoseconds), `-1` if there is no tasks in the corresponding dispatcher at all.
+ *
+ * ### Invariants
+ *
+ *  - When invoked from [Dispatchers.Default] **thread** (even if the actual context is different dispatcher,
+ *    [CoroutineDispatcher.limitedParallelism] or any in-place wrapper), it runs an arbitrary task that ended
+ *    up being scheduled to [Dispatchers.Default] or its counterpart. Tasks scheduled to [Dispatchers.IO]
+ *    **are not** executed[1].
+ *  - When invoked from [Dispatchers.IO] thread, the same rules apply, but for blocking tasks only.
+ *
+ * [1] -- this is purely technical limitation: the scheduler does not have "notify me when CPU token is available" API,
+ * and we cannot leave this method without leaving thread in its original state.
+ *
+ * ### Rationale
+ *
+ * This is an internal API that is intended to replace IDEA's core FJP decomposition.
+ * The following API is provided by IDEA core:
+ * ```
+ * runDecomposedTaskAndJoinIt { // <- non-suspending call
+ *     // spawn as many task as needed
+ *     // these tasks can also invoke 'runDecomposedTaskAndJoinIt'
+ * }
+ * ```
+ * The key observation here is that 'runDecomposedTaskAndJoinIt' can be invoked from `Dispatchers.Default` itself,
+ * thus blocking at least one thread. To avoid deadlocks and starvation during large hierarchical decompositions,
+ * 'runDecomposedTaskAndJoinIt' should not only block, but also **help** executing the task or other tasks
+ * until an arbitrary condition is satisfied.
+ *
+ * See #3439 for additional details.
+ *
+ * ### Limitations and caveats
+ *
+ * - Executes tasks in-place, thus potentially leaking irrelevant thread-locals from the current thread
+ * - Is not 100% effective, it waits for [Long] returned nanoseconds even when work arrives immediately after returning
+ *   from this method
+ * - When there is no more work, it's up to the caller to decide what to do. It's important to account that
+ *   work to current dispatcher may arrive **later** from external sources [1]
+ *
+ * [1] -- this is also a technicality that can be solved in kotlinx.coroutines itself, but unfortunately requires
+ *        a tremendous effort.
+ *
+ * @throws IllegalStateException if the current thread is not system dispatcher thread
+ */
+@InternalCoroutinesApi
+@DelicateCoroutinesApi
+public fun runSingleTaskFromCurrentSystemDispatcher(): Long {
+    val thread = Thread.currentThread()
+    if (thread !is CoroutineScheduler.Worker) throw IllegalStateException("Expected CoroutineScheduler.Worker, but got $thread")
+    return thread.runSingleTask()
+}
diff --git a/kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt b/kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt
@@ -720,6 +720,28 @@ internal class CoroutineScheduler(
             tryReleaseCpu(WorkerState.TERMINATED)
         }
 
+        /**
+         * See [runSingleTaskFromCurrentSystemDispatcher] for rationale and details.
+         * This is a fine-tailored method for a specific use-case not expected to be used widely.
+         */
+        fun runSingleTask(): Long {
+            val stateSnapshot = state
+            val isCpuThread  = state == WorkerState.CPU_ACQUIRED
+            val task = if (isCpuThread) {
+                findCpuTask()
+            } else {
+                findBlockingTask()
+            }
+            if (task == null) {
+                if (minDelayUntilStealableTaskNs == 0L) return -1L
+                return minDelayUntilStealableTaskNs
+            }
+            runSafely(task)
+            if (!isCpuThread) decrementBlockingTasks()
+            assert { state == stateSnapshot}
+            return 0L
+        }
+
         // Counterpart to "tryUnpark"
         private fun tryPark() {
             if (!inStack()) {
@@ -879,9 +901,19 @@ internal class CoroutineScheduler(
              * * Check if our queue has one (maybe mixed in with CPU tasks)
              * * Poll global and try steal
              */
+            return findBlockingTask()
+        }
+
+        private fun findBlockingTask(): Task? {
             return localQueue.pollBlocking()
                 ?: globalBlockingQueue.removeFirstOrNull()
-                ?: trySteal(blockingOnly = true)
+                ?: trySteal(STEAL_BLOCKING_ONLY)
+        }
+
+        private fun findCpuTask(): Task? {
+            return localQueue.pollCpu()
+                ?: globalBlockingQueue.removeFirstOrNull()
+                ?: trySteal(STEAL_CPU_ONLY)
         }
 
         private fun findAnyTask(scanLocalQueue: Boolean): Task? {
@@ -897,7 +929,7 @@ internal class CoroutineScheduler(
             } else {
                 pollGlobalQueues()?.let { return it }
             }
-            return trySteal(blockingOnly = false)
+            return trySteal(STEAL_ANY)
         }
 
         private fun pollGlobalQueues(): Task? {
@@ -910,7 +942,7 @@ internal class CoroutineScheduler(
             }
         }
 
-        private fun trySteal(blockingOnly: Boolean): Task? {
+        private fun trySteal(stealingMode: StealingMode): Task? {
             val created = createdWorkers
             // 0 to await an initialization and 1 to avoid excess stealing on single-core machines
             if (created < 2) {
@@ -924,11 +956,7 @@ internal class CoroutineScheduler(
                 if (currentIndex > created) currentIndex = 1
                 val worker = workers[currentIndex]
                 if (worker !== null && worker !== this) {
-                    val stealResult = if (blockingOnly) {
-                        worker.localQueue.tryStealBlocking(stolenTask)
-                    } else {
-                        worker.localQueue.trySteal(stolenTask)
-                    }
+                    val stealResult = worker.localQueue.trySteal(stealingMode, stolenTask)
                     if (stealResult == TASK_STOLEN) {
                         val result = stolenTask.element
                         stolenTask.element = null
diff --git a/kotlinx-coroutines-core/jvm/src/scheduling/WorkQueue.kt b/kotlinx-coroutines-core/jvm/src/scheduling/WorkQueue.kt
@@ -16,6 +16,11 @@ internal const val MASK = BUFFER_CAPACITY - 1 // 128 by default
 internal const val TASK_STOLEN = -1L
 internal const val NOTHING_TO_STEAL = -2L
 
+internal typealias StealingMode = Int
+internal const val STEAL_ANY: StealingMode = -1
+internal const val STEAL_CPU_ONLY: StealingMode = 0
+internal const val STEAL_BLOCKING_ONLY: StealingMode = 1
+
 /**
  * Tightly coupled with [CoroutineScheduler] queue of pending tasks, but extracted to separate file for simplicity.
  * At any moment queue is used only by [CoroutineScheduler.Worker] threads, has only one producer (worker owning this queue)
@@ -108,24 +113,34 @@ internal class WorkQueue {
      * Returns [NOTHING_TO_STEAL] if queue has nothing to steal, [TASK_STOLEN] if at least task was stolen
      * or positive value of how many nanoseconds should pass until the head of this queue will be available to steal.
      */
-    fun trySteal(stolenTaskRef: ObjectRef<Task?>): Long {
-        val task  = pollBuffer()
+    // TODO move it to tests where appropriate
+    fun trySteal(stolenTaskRef: ObjectRef<Task?>): Long = trySteal(STEAL_ANY, stolenTaskRef)
+
+    fun trySteal(stealingMode: StealingMode, stolenTaskRef: ObjectRef<Task?>): Long {
+        val task = when (stealingMode) {
+            STEAL_ANY -> pollBuffer()
+            else -> stealWithExclusiveMode(stealingMode)
+        }
+
         if (task != null) {
             stolenTaskRef.element = task
             return TASK_STOLEN
         }
-        return tryStealLastScheduled(stolenTaskRef, blockingOnly = false)
+        return tryStealLastScheduled(stealingMode, stolenTaskRef)
     }
 
-    fun tryStealBlocking(stolenTaskRef: ObjectRef<Task?>): Long {
+    // Steal only tasks of a particular kind, potentially invoking full queue scan
+    private fun stealWithExclusiveMode(stealingMode: StealingMode): Task? {
         var start = consumerIndex.value
         val end = producerIndex.value
-
-        while (start != end && blockingTasksInBuffer.value > 0) {
-            stolenTaskRef.element = tryExtractBlockingTask(start++) ?: continue
-            return TASK_STOLEN
+        val onlyBlocking = stealingMode == STEAL_BLOCKING_ONLY
+        // CPU or (BLOCKING & hasBlocking)
+        val shouldProceed = !onlyBlocking || blockingTasksInBuffer.value > 0
+        while (start != end && shouldProceed) {
+            return tryExtractFromTheMiddle(start++, onlyBlocking) ?: continue
         }
-        return tryStealLastScheduled(stolenTaskRef, blockingOnly = true)
+
+        return null
     }
 
     // Polls for blocking task, invoked only by the owner
@@ -138,23 +153,41 @@ internal class WorkQueue {
             } // Failed -> someone else stole it
         }
 
+        return pollWithMode(onlyBlocking = true /* only blocking */)
+    }
+
+    fun pollCpu(): 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
+        }
+
+        return pollWithMode(onlyBlocking = false /* only cpu */)
+    }
+
+    private fun pollWithMode(/* Only blocking OR only CPU */onlyBlocking: Boolean): Task? {
         val start = consumerIndex.value
         var end = producerIndex.value
-
-        while (start != end && blockingTasksInBuffer.value > 0) {
-            val task = tryExtractBlockingTask(--end)
+        // CPU or (BLOCKING & hasBlocking)
+        val shouldProceed = !onlyBlocking || blockingTasksInBuffer.value > 0
+        while (start != end && shouldProceed) {
+            val task = tryExtractFromTheMiddle(--end, onlyBlocking)
             if (task != null) {
                 return task
             }
         }
         return null
     }
 
-    private fun tryExtractBlockingTask(index: Int): Task? {
+    private fun tryExtractFromTheMiddle(index: Int, onlyBlocking: Boolean): Task? {
+        if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
         val arrayIndex = index and MASK
         val value = buffer[arrayIndex]
-        if (value != null && value.isBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
-            blockingTasksInBuffer.decrementAndGet()
+        if (value != null && value.isBlocking == onlyBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
+            if (onlyBlocking) blockingTasksInBuffer.decrementAndGet()
             return value
         }
         return null
@@ -170,10 +203,16 @@ internal class WorkQueue {
     /**
      * Contract on return value is the same as for [trySteal]
      */
-    private fun tryStealLastScheduled(stolenTaskRef: ObjectRef<Task?>, blockingOnly: Boolean): Long {
+    private fun tryStealLastScheduled(stealingMode: StealingMode, stolenTaskRef: ObjectRef<Task?>): Long {
         while (true) {
             val lastScheduled = lastScheduledTask.value ?: return NOTHING_TO_STEAL
-            if (blockingOnly && !lastScheduled.isBlocking) return NOTHING_TO_STEAL
+            if (lastScheduled.isBlocking) {
+                if (stealingMode == STEAL_CPU_ONLY) {
+                    return NOTHING_TO_STEAL
+                }
+            } else if (stealingMode == STEAL_BLOCKING_ONLY) {
+                return NOTHING_TO_STEAL
+            }
 
             // TODO time wraparound ?
             val time = schedulerTimeSource.nanoTime()
