IJPL-177878 Implement generic function for parallelism compensation

Author: knisht

IntelliJ-patches.md

@@ -45,11 +45,19 @@ To acquire an analogue of `limitedParallelism` dispatcher which supports paralle
 with `.softLimitedParallelism`, e.g., `.limitedParallelism(1)` may be used as a synchronization manager and in this case
 exceeding the parallelism limit would eliminate this (likely expected) side effect.
 
+### General parallelism compensation
+
+`runBlocking` is not the only operation that causes thread starvation. In the IntelliJ platform, we are also dealing with
+ the acquisition of read lock, which is essentially a blocking operation on a monitor. Any code can block on this monitor, 
+which results in a deadlock when there are no threads in a limited-thread dispatcher that are able to cancel the cancelling coroutines.
+We provide a user-visible function for parallelism compensation for arbitrary blocking operation.
+
 ### API
 - `runBlockingWithParallelismCompensation` - an analogue of `runBlocking` which also compensates parallelism of the
   associated coroutine dispatcher when it decides to park the thread
 - `CoroutineDispatcher.softLimitedParallelism` – an analogue of `.limitedParallelism` which supports
   parallelism compensation
+- `runAndCompensateParallelism` - a wrapper for an arbitrary blocking operation that initiates parallelism compensation after deadline.
 
 ## Asynchronous stack traces for flows in the IDEA debugger
 

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

@@ -7,6 +7,7 @@ package kotlinx.coroutines
 import kotlinx.coroutines.scheduling.withCompensatedParallelism
 import kotlin.contracts.*
 import kotlin.coroutines.*
+import kotlin.time.Duration
 
 /**
  * Runs a new coroutine and **blocks** the current thread _interruptibly_ until its completion.
@@ -113,7 +114,7 @@ private class BlockingCoroutine<T>(
                     if (isCompleted) break
                     if (parkNanos > 0) {
                         if (compensateParallelism) {
-                            withCompensatedParallelism {
+                            withCompensatedParallelism(Duration.ZERO) {
                                 parkNanos(this, parkNanos)
                             }
                         } else {

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

@@ -11,8 +11,10 @@ import kotlinx.coroutines.internal.softLimitedParallelism as softLimitedParallel
 import kotlinx.coroutines.internal.SoftLimitedDispatcher
 import kotlinx.coroutines.runBlockingWithParallelismCompensation as runBlockingWithParallelismCompensationImpl
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.scheduling.withCompensatedParallelism
 import kotlin.coroutines.CoroutineContext
 import kotlin.jvm.Throws
+import kotlin.time.Duration
 
 internal val currentContextThreadLocal : ThreadLocal<CoroutineContext?> = ThreadLocal.withInitial { null }
 
@@ -60,4 +62,23 @@ public object IntellijCoroutines {
     @Deprecated("use named version", level = DeprecationLevel.HIDDEN)
     public fun CoroutineDispatcher.softLimitedParallelism(parallelism: Int): CoroutineDispatcher =
         softLimitedParallelismImpl(parallelism, null)
+
+    /**
+     * Executes [action] and **advises** to compensate parallelism if [action] does not finish within [timeout].
+     *
+     * Suppose that you are dealing with an operation that blocks the underlying thread.
+     * If the blocked thread happens to be from a limited-thread-pool dispatcher, then you face the problem of thread starvation:
+     * your system underutilizes the available CPU. In severe cases, this may result in a deadlock.
+     * To increase the size of the thread pool, you can wrap the operation into this function.
+     *
+     * Remember that threads are expensive, so you need to choose an appropriate [timeout]
+     * so that parallelism compensation would be less expensive than waiting for [action] to finish.
+     *
+     * There are intentionally few guarantees that this function will have side effects.
+     * Generally, parallelism compensation applies only to threads that are ready for compensation,
+     * such as [Dispatchers.Default] or [Dispatchers.IO].
+     */
+    public fun <T> runAndCompensateParallelism(timeout: Duration, action: () -> T): T {
+        return withCompensatedParallelism(timeout, action)
+    }
 }

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

@@ -1,5 +1,11 @@
 package kotlinx.coroutines.scheduling
 
+import kotlinx.coroutines.DefaultDelay
+import kotlinx.coroutines.DisposableHandle
+import kotlinx.coroutines.internal.synchronized
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.time.Duration
+
 private val parallelismCompensationEnabled: Boolean =
     System.getProperty("kotlinx.coroutines.parallelism.compensation", "true").toBoolean()
 
@@ -11,21 +17,80 @@ private val parallelismCompensationEnabled: Boolean =
  * that eventually it will adjust to meet it.
  */
 @Suppress("NOTHING_TO_INLINE") // better stacktrace
-internal inline fun <T> withCompensatedParallelism(noinline body: () -> T): T {
+internal inline fun <T> withCompensatedParallelism(timeout: Duration, noinline body: () -> T): T {
     if (!parallelismCompensationEnabled) {
         return body()
     }
     // CoroutineScheduler.Worker implements ParallelismCompensation
     val worker = Thread.currentThread() as? ParallelismCompensation
         ?: return body()
-    return worker.withCompensatedParallelism(body)
+    return if (timeout == Duration.ZERO) {
+        worker.withCompensatedParallelismImmediate(body)
+    } else {
+        worker.withCompensatedParallelismAfterDeadline(timeout, body)
+    }
 }
 
-private fun <T> ParallelismCompensation.withCompensatedParallelism(body: () -> T): T {
+private fun <T> ParallelismCompensation.withCompensatedParallelismImmediate(body: () -> T): T {
     increaseParallelismAndLimit()
     try {
         return body()
     } finally {
         decreaseParallelismLimit()
     }
+}
+
+private fun <T> ParallelismCompensation.withCompensatedParallelismAfterDeadline(timeout: Duration, body: () -> T): T {
+    val holder = WorkerProtector()
+    val compensationHandle = initCompensationWithDeadline(this, holder, timeout)
+    try {
+        return body()
+    } finally {
+        endCompensationWithDeadline(this, holder, compensationHandle)
+    }
+}
+
+/**
+ * The fields in [CoroutineScheduler.Worker] are not thread safe: these are plain java fields, and they are intended to
+ * be accessed only by [CoroutineScheduler.Worker] itself. Since we want to modify them from other threads,
+ * we need to establish proper happens-before relations for accesses to these fields.
+ */
+private class WorkerProtector() {
+    @Volatile
+    var wasTaken: Boolean = false
+}
+
+private fun initCompensationWithDeadline(
+    thread: ParallelismCompensation,
+    state: WorkerProtector,
+    timeout: Duration
+): DisposableHandle {
+    return DefaultDelay.invokeOnTimeout(timeout.inWholeMilliseconds, {
+        if (!state.wasTaken) {
+            synchronized(state) {
+                if (!state.wasTaken) {
+                    // so the task was not finished yet, which means that we can increase parallelism here
+                    thread.increaseParallelismAndLimit()
+                    state.wasTaken = true
+                }
+            }
+        }
+    }, EmptyCoroutineContext)
+}
+
+private fun endCompensationWithDeadline(
+    thread: ParallelismCompensation,
+    state: WorkerProtector,
+    compensationHandle: DisposableHandle,
+) {
+    synchronized(state) {
+        if (!state.wasTaken) {
+            // so parallelism was not compensated at this moment; we can abort the compensating coroutine
+            compensationHandle.dispose()
+            state.wasTaken = true
+        } else {
+            // the compensating coroutine managed to publish a request for compensation. Now we need to decompensate
+            thread.decreaseParallelismLimit()
+        }
+    }
 }

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

@@ -0,0 +1,55 @@
+package kotlinx.coroutines.scheduling
+
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.internal.intellij.IntellijCoroutines
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import java.util.concurrent.CountDownLatch
+import kotlin.test.assertTrue
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
+
+class GeneralParallelismCompensationTest : SchedulerTestBase() {
+
+    @Test
+    fun `soundness - parallelism compensation ensures progress of operations`() {
+        val operationsCount = CORE_POOL_SIZE * 10
+        val barrier = CountDownLatch(operationsCount)
+        runBlocking(Dispatchers.Default) {
+            repeat(operationsCount) {
+                launch {
+                    IntellijCoroutines.runAndCompensateParallelism(10.milliseconds) {
+                        barrier.countDown()
+                        barrier.await()
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    fun `laziness - parallelism compensation does not do anything before deadline`() {
+        val operationsCount = CORE_POOL_SIZE * 3
+        val barrier = CountDownLatch(operationsCount)
+        runBlocking(Dispatchers.Default) {
+            launch(Dispatchers.IO) {
+                delay(1.seconds)
+                assertTrue(barrier.count == CORE_POOL_SIZE.toLong() * 2)
+                delay(2.seconds)
+                assertTrue(barrier.count == CORE_POOL_SIZE.toLong())
+                delay(2.seconds)
+                assertTrue(barrier.count == 0L)
+            }
+            repeat(operationsCount) {
+                launch {
+                    IntellijCoroutines.runAndCompensateParallelism(2.seconds) {
+                        barrier.countDown()
+                        barrier.await()
+                    }
+                }
+            }
+        }
+    }
+}