Introduce withUnlimitedIOScheduler utility method to use instead of direct interaction with Worker

Kotlin#3983 / IJPL-721

Author: vsalavatov

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

@@ -5,7 +5,6 @@
 package kotlinx.coroutines
 
 import kotlinx.coroutines.scheduling.*
-import kotlinx.coroutines.scheduling.CoroutineScheduler
 import java.util.concurrent.locks.*
 import kotlin.contracts.*
 import kotlin.coroutines.*
@@ -88,7 +87,6 @@ private class BlockingCoroutine<T>(
     @Suppress("UNCHECKED_CAST")
     fun joinBlocking(): T {
         registerTimeLoopThread()
-        var cpuPermitReleased: Boolean? = null
         try {
             eventLoop?.incrementUseCount()
             try {
@@ -98,20 +96,17 @@ 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
-                    if (parkNanos > 0 && cpuPermitReleased == null) {
-                        val worker = Thread.currentThread() as? CoroutineScheduler.Worker
-                        cpuPermitReleased = worker?.releaseCpu() ?: false
+                    if (parkNanos > 0) {
+                        withUnlimitedIOScheduler {
+                            parkNanos(this, parkNanos)
+                        }
                     }
-                    parkNanos(this, parkNanos)
                 }
             } finally { // paranoia
                 eventLoop?.decrementUseCount()
             }
         } finally { // paranoia
             unregisterTimeLoopThread()
-            if (cpuPermitReleased == true) {
-                (Thread.currentThread() as CoroutineScheduler.Worker).reacquireCpu()
-            }
         }
         // now return result
         val state = this.state.unboxState()

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

@@ -690,15 +690,15 @@ internal class CoroutineScheduler(
             return hadCpu
         }
 
-        /** only for [runBlocking] */
+        /** only for [withoutCpuPermit] */
         fun releaseCpu(): Boolean {
             assert { state == WorkerState.CPU_ACQUIRED || state == WorkerState.BLOCKING }
             return tryReleaseCpu(WorkerState.BLOCKING).also { released ->
                 if (released) incrementBlockingTasks()
             }
         }
 
-        /** only for [runBlocking] */
+        /** only for [withoutCpuPermit] */
         fun reacquireCpu() {
             assert { state == WorkerState.BLOCKING }
             decrementBlockingTasks()
@@ -732,7 +732,7 @@ internal class CoroutineScheduler(
                     assert { state == WorkerState.CPU_ACQUIRED }
                     break
                 }
-                LockSupport.parkNanos(RUN_BLOCKING_CPU_REACQUIRE_PARK_NS)
+                LockSupport.parkNanos(CPU_REACQUIRE_PARK_NS)
             }
         }
 
@@ -1085,3 +1085,22 @@ internal fun isSchedulerWorker(thread: Thread) = thread is CoroutineScheduler.Wo
 @JvmName("mayNotBlock")
 internal fun mayNotBlock(thread: Thread) = thread is CoroutineScheduler.Worker &&
     thread.state == CoroutineScheduler.WorkerState.CPU_ACQUIRED
+
+/**
+ * Emulates dispatch to [UnlimitedIoScheduler] in a blocking context.
+ */
+internal fun withUnlimitedIOScheduler(blocking: () -> Unit) {
+    withoutCpuPermit {
+        blocking()
+    }
+}
+
+private fun withoutCpuPermit(body: () -> Unit) {
+    val worker = Thread.currentThread() as? CoroutineScheduler.Worker ?: return body()
+    val releasedPermit = worker.releaseCpu()
+    try {
+        return body()
+    } finally {
+        if (releasedPermit) worker.reacquireCpu()
+    }
+}

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

@@ -20,8 +20,8 @@ internal val WORK_STEALING_TIME_RESOLUTION_NS = systemProp(
 )
 
 @JvmField
-internal val RUN_BLOCKING_CPU_REACQUIRE_PARK_NS = systemProp(
-    "kotlinx.coroutines.scheduler.runBlocking.cpu.reacquire.ns", 250L * 1000 * 1000
+internal val CPU_REACQUIRE_PARK_NS = systemProp(
+    "kotlinx.coroutines.scheduler.cpu.reacquire.ns", 250L * 1000 * 1000
 )
 
 /**