Avoid potential StackOverflow on wrapped UnconfinedDispatcher+yield

See ImmediateYieldTest.testWrappedUnconfinedDispatcherYieldStackOverflow

This commit changes the way Unconfined dispatcher is detected.
Instead of equality check we try to dispatch to the dispatcher with
a specially updated coroutine context that has YieldContext element
that is processed by the Unconfied dispatcher in a special way.

kotlinx-coroutines-core/common/src/CoroutineDispatcher.kt

@@ -76,6 +76,10 @@ public abstract class CoroutineDispatcher :
  *
  * This method should generally be exception-safe. An exception thrown from this method
  * may leave the coroutines that use this dispatcher in the inconsistent and hard to debug state.
+ *
+ * **Note**: This method must not immediately call [block]. Doing so would result in [StackOverflowError]
+ * when [yield] is repeatedly called from a loop. However, an implementation can delegate this function
+ * to `dispatch` method of [Dispatchers.Unconfined], which is integrated with [yield] to avoid this problem.
  */
  public abstract fun dispatch(context: CoroutineContext, block: Runnable)
 

kotlinx-coroutines-core/common/src/Unconfined.kt

@@ -5,13 +5,35 @@
 package kotlinx.coroutines
 
 import kotlin.coroutines.*
+import kotlin.jvm.*
 
 /**
  * A coroutine dispatcher that is not confined to any specific thread.
  */
 internal object Unconfined : CoroutineDispatcher() {
  override fun isDispatchNeeded(context: CoroutineContext): Boolean = false
- // Just in case somebody wraps Unconfined dispatcher casing the "dispatch" to be called from "yield"
- override fun dispatch(context: CoroutineContext, block: Runnable) = block.run()
+
+ override fun dispatch(context: CoroutineContext, block: Runnable) {
+ // Just in case somebody wraps Unconfined dispatcher casing the "dispatch" to be called from "yield"
+ // See also code of "yield" function
+ val yieldContext = context[YieldContext]
+ if (yieldContext != null) {
+ // report to "yield" that it is an unconfined dispatcher and don't call "block.run()"
+ yieldContext.dispatcherWasUnconfined = true
+ return
+ }
+ block.run()
+ }
+
  override fun toString(): String = "Unconfined"
 }
+
+/**
+ * Used to detect calls to [Unconfined.dispatch] from [yield] function.
+ */
+internal class YieldContext : AbstractCoroutineContextElement(Key) {
+ companion object Key : CoroutineContext.Key<YieldContext>
+
+ @JvmField
+ var dispatcherWasUnconfined = false
+}

kotlinx-coroutines-core/common/src/Yield.kt

@@ -27,9 +27,16 @@ public suspend fun yield(): Unit = suspendCoroutineUninterceptedOrReturn sc@ { u
  val context = uCont.context
  context.checkCompletion()
  val cont = uCont.intercepted() as? DispatchedContinuation<Unit> ?: return@sc Unit
+ // This code detects the Unconfined dispatcher even if it was wrapped into another dispatcher
+ val yieldContext = YieldContext()
+ cont.dispatchYield(context + yieldContext, Unit)
  // Special case for the unconfined dispatcher that can yield only in existing unconfined loop
- if (cont.dispatcher === Unconfined) return@sc if (cont.yieldUndispatched()) COROUTINE_SUSPENDED else Unit
- cont.dispatchYield(Unit)
+ if (yieldContext.dispatcherWasUnconfined) {
+ // Means that the Unconfined dispatcher got the call, but did not do anything.
+ // See also code of "Unconfined.dispatch" function.
+ return@sc if (cont.yieldUndispatched()) COROUTINE_SUSPENDED else Unit
+ }
+ // It was some other dispatcher that successfully dispatched the coroutine
  COROUTINE_SUSPENDED
 }
 

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

@@ -211,8 +211,7 @@ internal class DispatchedContinuation<in T>(
  }
 
  // used by "yield" implementation
- internal fun dispatchYield(value: T) {
- val context = continuation.context
+ internal fun dispatchYield(context: CoroutineContext, value: T) {
  _state = value
  resumeMode = MODE_CANCELLABLE
  dispatcher.dispatchYield(context, this)

kotlinx-coroutines-core/common/test/ImmediateYieldTest.kt

@@ -43,4 +43,15 @@ class ImmediateYieldTest : TestBase() {
  }
  finish(4) // after launch
  }
+
+ @Test
+ fun testWrappedUnconfinedDispatcherYieldStackOverflow() = runTest {
+ expect(1)
+ withContext(wrapperDispatcher(Dispatchers.Unconfined)) {
+ repeat(100_000) {
+ yield()
+ }
+ }
+ finish(2)
+ }
 }