Make runBlocking release worker permits before park and reacquire after unpark to avoid starvation

IntelliJ-patches.md

@@ -0,0 +1,20 @@
+# Included IntelliJ-related patches
+
+## `runBlocking` without Dispatcher starvation 
+
+[IJPL-721](https://youtrack.jetbrains.com/issue/IJPL-721), [#3983](https://github.com/Kotlin/kotlinx.coroutines/issues/3983)
+
+### Description:
+`runBlocking` with its default semantics may cause dispatcher starvation if it is called on a worker thread. 
+For example, if `runBlocking` happens to block all `Dispatchers.Default` workers, it may lead to a deadlock in the application:
+there may be tasks in the CPU queue that `runBlocking`s await, but there are none CPU workers available to run them.
+
+This patch changes the behavior of `runBlocking` so that it always releases associated computation permits before it parks,
+and reacquires them after unpark. It works for every `CoroutineDispatcher` that is built using library primitives:
+plain `Dispatcher.*` objects or `.limitedParallelism` limited dispatchers that are on top of them.
+
+This change in behavior comes with a cost. Permit reacquiring mechanism _may_ need an additional thread park/unpark.
+Worker threads that release their computational permits always let go of the local task queue, which means less benefit
+from locality, higher contention and transactional costs at the very least. 
+
+This patch doesn't change the fact that `runBlocking` should still be used carefully.

README.md

@@ -1,4 +1,14 @@
-# kotlinx.coroutines 
+# kotlinx.coroutines with IntelliJ patches
+
+This repository is a fork of the original [kotlinx.coroutines](https://github.com/Kotlin/kotlinx.coroutines) library that includes
+several patches (see [IntelliJ-patches.md](IntelliJ-patches.md)).
+
+**Important:** 
+- Only JVM build target is fully supported and used
+
+Release instructions are [here](RELEASE.md).
+
+---
 
 [![Kotlin Stable](https://kotl.in/badges/stable.svg)](https://kotlinlang.org/docs/components-stability.html)
 [![JetBrains official project](https://jb.gg/badges/official.svg)](https://confluence.jetbrains.com/display/ALL/JetBrains+on+GitHub)

RELEASE.md

@@ -1,3 +1,25 @@
+# com.intellij.platform:kotlinx-coroutines-* release
+
+```
+# update branches
+# git remote add upstream https://github.com/Kotlin/kotlinx.coroutines.git 
+git checkout master
+git fetch upstream
+git fetch origin
+
+# prepare new master with patches
+git reset --hard upstream/master
+git rebase intellij/patch-base
+git rebase intellij/whatever-patches-we-have, see IntelliJ-patches.md for the list of branches
+
+# Remember to change the version in `gradle.properties` to something like `1.8.4-intellij-SNAPSHOT`
+# commit version change
+
+git push origin master --force
+```
+
+---
+
 # kotlinx.coroutines release checklist
 
 To release a new `<version>` of `kotlinx-coroutines`:

buildSrc/src/main/kotlin/Publishing.kt

@@ -48,11 +48,13 @@ fun MavenPom.configureMavenCentralMetadata(project: Project) {
  * dev build into 'https://maven.pkg.jetbrains.space/public/p/kotlinx-coroutines/maven' Maven repository.
  * In order to use it, pass the corresponding ENV to the TC 'Deploy' task.
  */
-private val spacePublicationEnabled = System.getenv("libs.space.pub")?.equals("true") ?: false
+private val Project.spacePublicationEnabled: Boolean
+    get() = getSensitiveProperty("libs.space.pub")?.equals("true") ?: false
 
-fun mavenRepositoryUri(): URI {
+fun Project.mavenRepositoryUri(): URI {
     if (spacePublicationEnabled) {
-        return URI("https://maven.pkg.jetbrains.space/public/p/kotlinx-coroutines/maven")
+        val spaceRepoUrl = getSensitiveProperty("libs.space.url")
+        return URI(spaceRepoUrl ?: "https://maven.pkg.jetbrains.space/public/p/kotlinx-coroutines/maven")
     }
 
     val repositoryId: String? = System.getenv("libs.repository.id")
@@ -65,9 +67,9 @@ fun mavenRepositoryUri(): URI {
 
 fun configureMavenPublication(rh: RepositoryHandler, project: Project) {
     rh.maven {
-        url = mavenRepositoryUri()
+        url = project.mavenRepositoryUri()
         credentials {
-            if (spacePublicationEnabled) {
+            if (project.spacePublicationEnabled) {
                 // Configure space credentials
                 username = project.getSensitiveProperty("libs.space.user")
                 password = project.getSensitiveProperty("libs.space.password")

gradle.properties

@@ -1,6 +1,6 @@
 # Kotlin
 version=1.8.0-SNAPSHOT
-group=org.jetbrains.kotlinx
+group=com.intellij.platform
 kotlin_version=1.9.21
 # DO NOT rename this property without adapting kotlinx.train build chain:
 atomicfu_version=0.23.1

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

@@ -0,0 +1,22 @@
+package kotlinx.coroutines.internal
+
+/**
+ * [Runnables][kotlinx.coroutines.Runnable] that are dispatched on [CoroutineDispatcher][kotlinx.coroutines.CoroutineDispatcher]
+ * may optionally implement this interface to be notified of dispatch emulation in blocking mode.
+ *
+ * This may be needed for controlling dispatcher to release/acquire a permit of the worker that currently
+ * executes the dispatched Runnable.
+ * @see LimitedDispatcher.Worker
+ * @see kotlinx.coroutines.scheduling.withUnlimitedIOScheduler
+ */
+internal interface BlockingDispatchAware {
+    /**
+     * Must not throw
+     */
+    fun beforeDispatchElsewhere()
+
+    /**
+     * Must not throw
+     */
+    fun afterDispatchBack()
+}

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

@@ -103,12 +103,14 @@ internal class LimitedDispatcher(
      * 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 {
+    private inner class Worker(private var currentTask: Runnable) : Runnable, BlockingDispatchAware {
         override fun run() {
             var fairnessCounter = 0
             while (true) {
                 try {
                     currentTask.run()
+                } catch (e: WorkerPermitTransferCompleted) {
+                    if (!tryAllocateWorker()) return
                 } catch (e: Throwable) {
                     handleCoroutineException(EmptyCoroutineContext, e)
                 }
@@ -122,7 +124,31 @@ internal class LimitedDispatcher(
                 }
             }
         }
+
+        override fun beforeDispatchElsewhere() {
+            // compensate while we are blocked and consider that we gave away our permit to a new worker
+            val newWorker = Worker(Runnable {})
+            dispatcher.dispatch(this@LimitedDispatcher, newWorker)
+            (currentTask as? BlockingDispatchAware)?.beforeDispatchElsewhere()
+        }
+
+        override fun afterDispatchBack() {
+            (currentTask as? BlockingDispatchAware)?.afterDispatchBack()
+            if (tryAllocateWorker()) return
+            val permitTransfer = PermitTransfer()
+            queue.addLast(
+                permitTransfer.releaseFun { throw WorkerPermitTransferCompleted }
+                    .let { Runnable { it() } }
+            )
+            permitTransfer.acquire(
+                tryAllocatePermit = ::tryAllocateWorker,
+                deallocatePermit = { runningWorkers.decrementAndGet() }
+            )
+        }
+
     }
+
+    private object WorkerPermitTransferCompleted : Throwable()
 }
 
 // Save a few bytecode ops

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

@@ -0,0 +1,45 @@
+package kotlinx.coroutines.internal
+
+import kotlinx.atomicfu.*
+
+internal class PermitTransferStatus {
+    private val status = atomic(false)
+    fun check(): Boolean = status.value
+    fun complete(): Boolean = status.compareAndSet(false, true)
+}
+
+internal expect class PermitTransfer constructor() {
+    /**
+     * [releasePermit] may throw
+     */
+    fun releaseFun(releasePermit: () -> Unit): () -> Unit
+
+    /**
+     * [tryAllocatePermit] and [deallocatePermit] must not throw
+     */
+    fun acquire(tryAllocatePermit: () -> Boolean, deallocatePermit: () -> Unit)
+}
+
+internal class BusyPermitTransfer {
+    private val status = PermitTransferStatus()
+
+    fun releaseFun(releasePermit: () -> Unit): () -> Unit = {
+        if (status.complete()) {
+            releasePermit()
+        }
+    }
+
+    fun acquire(tryAllocatePermit: () -> Boolean, deallocatePermit: () -> Unit) {
+        while (true) {
+            if (status.check()) {
+                return
+            }
+            if (tryAllocatePermit()) {
+                if (!status.complete()) { // race: transfer was completed first by releaseFun
+                    deallocatePermit()
+                }
+                return
+            }
+        }
+    }
+}

kotlinx-coroutines-core/jsAndWasmShared/src/internal/PermitTransfer.jsAndWasmShared.kt

@@ -0,0 +1,3 @@
+package kotlinx.coroutines.internal
+
+internal actual typealias PermitTransfer = BusyPermitTransfer // TODO

kotlinx-coroutines-core/jvm/src/Builders.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.scheduling.*
 import java.util.concurrent.locks.*
 import kotlin.contracts.*
 import kotlin.coroutines.*
@@ -95,7 +96,11 @@ private class BlockingCoroutine<T>(
                     val parkNanos = eventLoop?.processNextEvent() ?: Long.MAX_VALUE
                     // note: process next even may loose unpark flag, so check if completed before parking
                     if (isCompleted) break
-                    parkNanos(this, parkNanos)
+                    if (parkNanos > 0) {
+                        withUnlimitedIOScheduler {
+                            parkNanos(this, parkNanos)
+                        }
+                    }
                 }
             } finally { // paranoia
                 eventLoop?.decrementUseCount()

kotlinx-coroutines-core/jvm/src/internal/PermitTransfer.jvm.kt

@@ -0,0 +1,35 @@
+package kotlinx.coroutines.internal
+
+import java.util.concurrent.locks.*
+
+internal actual class PermitTransfer {
+    private val status = PermitTransferStatus()
+
+    public actual fun releaseFun(releasePermit: () -> Unit): () -> Unit {
+        val blockedThread = Thread.currentThread()
+        return {
+            if (status.complete()) {
+                try {
+                    releasePermit()
+                } finally {
+                    LockSupport.unpark(blockedThread)
+                }
+            }
+        }
+    }
+
+    public actual fun acquire(tryAllocatePermit: () -> Boolean, deallocatePermit: () -> Unit) {
+        while (true) {
+            if (status.check()) {
+                return
+            }
+            if (tryAllocatePermit()) {
+                if (!status.complete()) { // race: transfer was completed first by releaseFun
+                    deallocatePermit()
+                }
+                return
+            }
+            LockSupport.park(this)
+        }
+    }
+}