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.

Author: elizarov

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

@@ -133,7 +133,7 @@ public abstract class AbstractCoroutine<in T>(
      *
      * * [DEFAULT] uses [startCoroutineCancellable].
      * * [ATOMIC] uses [startCoroutine].
-     * * [UNDISPATCHED] uses [startCoroutineUndispatched].
+     * * [UNCONFINED] uses [startCoroutineUndispatched].
      * * [LAZY] does nothing.
      */
     public fun start(start: CoroutineStart, block: suspend () -> T) {
@@ -150,7 +150,7 @@ public abstract class AbstractCoroutine<in T>(
      * 
      * * [DEFAULT] uses [startCoroutineCancellable].
      * * [ATOMIC] uses [startCoroutine].
-     * * [UNDISPATCHED] uses [startCoroutineUndispatched].
+     * * [UNCONFINED] uses [startCoroutineUndispatched].
      * * [LAZY] does nothing.
      */
     public fun <R> start(start: CoroutineStart, receiver: R, block: suspend R.() -> T) {

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/CoroutineStart.kt

@@ -27,7 +27,7 @@ public enum class CoroutineStart {
      * invoked the coroutine builder continues execution.
      *
      * Note that [Dispatchers.Unconfined] always returns `false` from its [CoroutineDispatcher.isDispatchNeeded]
-     * function, so starting a coroutine with [Dispatchers.Unconfined] by [DEFAULT] is the same as using [UNDISPATCHED].
+     * function, so starting a coroutine with [Dispatchers.Unconfined] by [DEFAULT] is the same as using [UNCONFINED].
      *
      * If coroutine [Job] is cancelled before it even had a chance to start executing, then it will not start its
      * execution at all, but will complete with an exception.
@@ -79,7 +79,7 @@ public enum class CoroutineStart {
      *
      * * [DEFAULT] uses [startCoroutineCancellable].
      * * [ATOMIC] uses [startCoroutine].
-     * * [UNDISPATCHED] uses [startCoroutineUndispatched].
+     * * [UNCONFINED] uses [startCoroutineUndispatched].
      * * [LAZY] does nothing.
      *
      * @suppress **This an internal API and should not be used from general code.**
@@ -98,7 +98,7 @@ public enum class CoroutineStart {
      *
      * * [DEFAULT] uses [startCoroutineCancellable].
      * * [ATOMIC] uses [startCoroutine].
-     * * [UNDISPATCHED] uses [startCoroutineUndispatched].
+     * * [UNCONFINED] uses [startCoroutineUndispatched].
      * * [LAZY] does nothing.
      *
      * @suppress **This an internal API and should not be used from general code.**

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)
+    }
 }