Implement permit release also for LimitedDispatcher

PermitTransfer is extracted to be used both in CoroutineScheduler and in LimitedDispatcher.
BlockingDispatchAware interface is introduced for LimitedDispatcher.Worker to be accounted by CoroutineScheduler.

Kotlin#3983 / IJPL-721

Author: vsalavatov

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

@@ -0,0 +1,7 @@
+package kotlinx.coroutines.internal
+
+
+internal interface BlockingDispatchAware {
+    fun beforeDispatchElsewhere()
+    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) {
+                    return
                 } catch (e: Throwable) {
                     handleCoroutineException(EmptyCoroutineContext, e)
                 }
@@ -122,7 +124,29 @@ internal class LimitedDispatcher(
                 }
             }
         }
+
+        override fun beforeDispatchElsewhere() {
+            // compensate while we are blocked
+            val newWorker = Worker(Runnable {})
+            dispatcher.dispatch(this@LimitedDispatcher, newWorker)
+        }
+
+        override fun 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,51 @@
+package kotlinx.coroutines.internal
+
+import kotlinx.atomicfu.*
+import kotlin.jvm.*
+
+@JvmField
+internal val PERMIT_ACQUIRE_PARK_NS = systemProp(
+    "kotlinx.coroutines.permit.acquire.ns", 250L * 1000 * 1000 // 250ms
+)
+
+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/internal/PermitTransfer.jvm.kt

@@ -0,0 +1,37 @@
+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.parkNanos(
+                PERMIT_ACQUIRE_PARK_NS // todo: not sure if it's needed at all, I mean that it is < Long.MAX_VALUE, but at least this way it's safer
+            )
+        }
+    }
+}

kotlinx-coroutines-core/jvm/src/scheduling/CoroutineScheduler.kt

@@ -703,37 +703,18 @@ internal class CoroutineScheduler(
             assert { state == WorkerState.BLOCKING }
             decrementBlockingTasks()
             if (tryAcquireCpuPermit()) return
-            class CpuPermitTransfer {
-                private val status = atomic(false)
-                fun check(): Boolean = status.value
-                fun complete(): Boolean = status.compareAndSet(false, true)
-            }
-            val permitTransfer = CpuPermitTransfer()
-            val blockedWorker = this@Worker
-            scheduler.dispatch(Runnable {
+            val permitTransfer = PermitTransfer()
+            scheduler.dispatch(permitTransfer.releaseFun {
                 // this code runs in a different worker thread that holds a CPU token
                 val cpuHolder = currentThread() as Worker
                 assert { cpuHolder.state == WorkerState.CPU_ACQUIRED }
-                if (permitTransfer.complete()) {
-                    cpuHolder.state = WorkerState.BLOCKING
-                    LockSupport.unpark(blockedWorker)
-                }
+                cpuHolder.state = WorkerState.BLOCKING
             }, taskContext = NonBlockingContext)
-            while (true) {
-                if (permitTransfer.check()) {
-                    state = WorkerState.CPU_ACQUIRED
-                    break
-                }
-                if (tryAcquireCpuPermit()) {
-                    if (!permitTransfer.complete()) {
-                        // race: transfer was completed by another thread
-                        releaseCpuPermit()
-                    }
-                    assert { state == WorkerState.CPU_ACQUIRED }
-                    break
-                }
-                LockSupport.parkNanos(CPU_REACQUIRE_PARK_NS)
-            }
+            permitTransfer.acquire(
+                tryAllocatePermit = this@CoroutineScheduler::tryAcquireCpuPermit,
+                deallocatePermit = ::releaseCpuPermit
+            )
+            state = WorkerState.CPU_ACQUIRED
         }
 
         override fun run() = runWorker()
@@ -841,14 +822,20 @@ internal class CoroutineScheduler(
 
         private fun inStack(): Boolean = nextParkedWorker !== NOT_IN_STACK
 
+        private var currentTask: Task? = null
+
         private fun executeTask(task: Task) {
             val taskMode = task.mode
             idleReset(taskMode)
             beforeTask(taskMode)
+            currentTask = task
             runSafely(task)
+            currentTask = null
             afterTask(taskMode)
         }
 
+        internal fun getCurrentTaskImpl(): TaskImpl? = currentTask as? TaskImpl
+
         private fun beforeTask(taskMode: Int) {
             if (taskMode == TASK_NON_BLOCKING) return
             // Always notify about new work when releasing CPU-permit to execute some blocking task
@@ -1091,7 +1078,9 @@ internal fun mayNotBlock(thread: Thread) = thread is CoroutineScheduler.Worker &
  */
 internal fun withUnlimitedIOScheduler(blocking: () -> Unit) {
     withoutCpuPermit {
-        blocking()
+        withTaskBlockingDispatch {
+            blocking()
+        }
     }
 }
 
@@ -1104,3 +1093,14 @@ private fun withoutCpuPermit(body: () -> Unit) {
         if (releasedPermit) worker.reacquireCpu()
     }
 }
+
+private fun withTaskBlockingDispatch(body: () -> Unit) {
+    val worker = Thread.currentThread() as? CoroutineScheduler.Worker ?: return body()
+    val dispatchAware = worker.getCurrentTaskImpl()?.block as? BlockingDispatchAware ?: return body()
+    dispatchAware.beforeDispatchElsewhere()
+    try {
+        return body()
+    } finally {
+        dispatchAware.afterDispatchBack()
+    }
+}

kotlinx-coroutines-core/jvm/src/scheduling/Tasks.kt

@@ -19,10 +19,6 @@ internal val WORK_STEALING_TIME_RESOLUTION_NS = systemProp(
     "kotlinx.coroutines.scheduler.resolution.ns", 100000L
 )
 
-@JvmField
-internal val CPU_REACQUIRE_PARK_NS = systemProp(
-    "kotlinx.coroutines.scheduler.cpu.reacquire.ns", 250L * 1000 * 1000
-)
 
 /**
  * The maximum number of threads allocated for CPU-bound tasks at the default set of dispatchers.

kotlinx-coroutines-core/jvm/test/scheduling/BlockingCoroutineDispatcherTest.kt

@@ -4,6 +4,7 @@ import kotlinx.coroutines.*
 import org.junit.*
 import org.junit.rules.*
 import java.util.concurrent.*
+import java.util.concurrent.atomic.*
 
 class BlockingCoroutineDispatcherTest : SchedulerTestBase() {
 
@@ -186,4 +187,33 @@ class BlockingCoroutineDispatcherTest : SchedulerTestBase() {
         body(1)
         checkPoolThreadsCreated(maxDepth..maxDepth + 1)
     }
+
+    @Test
+    fun testNoStarvationOfLimitedDispatcherWithRunBlocking() {
+        val taskCount = 5
+        val dispatcher = blockingDispatcher(2)
+        val barrier = CompletableDeferred<Unit>()
+        val count = AtomicInteger(0)
+        fun blockingCode() {
+            runBlocking {
+                count.incrementAndGet()
+                barrier.await()
+                count.decrementAndGet()
+            }
+        }
+        runBlocking {
+            repeat(taskCount) {
+                launch(dispatcher) {
+                    blockingCode()
+                }
+            }
+            while (count.get() != taskCount) {
+                Thread.sleep(1)
+            }
+            barrier.complete(Unit)
+            while (count.get() != 0) {
+                Thread.sleep(1)
+            }
+        }
+    }
 }

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

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