Optimize yield via dispatcher.isDispatchNeeded

Also, as it was before, throw UnsupportedOperationException from the
Dispatcher.Unconfined.dispatch method in case some code wraps
the Unconfined dispatcher but fails to delegate isDispatchNeeded
properly.

Author: elizarov

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

@@ -78,8 +78,9 @@ public abstract class CoroutineDispatcher :
      * 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.
+     * when [yield] is repeatedly called from a loop. However, an implementation that returns `false` from
+     * [isDispatchNeeded] 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

@@ -14,15 +14,16 @@ internal object Unconfined : CoroutineDispatcher() {
     override fun isDispatchNeeded(context: CoroutineContext): Boolean = false
 
     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
+        // It can only be called by the "yield" function. 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()
+        throw UnsupportedOperationException("Dispatchers.Unconfined.dispatch function can only be used by the yield function. " +
+            "If you wrap Unconfined dispatcher in your code, make sure you properly delegate " +
+            "isDispatchNeeded and dispatch calls.")
     }
 
     override fun toString(): String = "Unconfined"

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

@@ -27,16 +27,22 @@ 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 (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
+    if (cont.dispatcher.isDispatchNeeded(context)) {
+        // this is a regular dispatcher -- do simple dispatchYield
+        cont.dispatchYield(context, Unit)
+    } else {
+        // This is either an "immediate" dispatcher or the Unconfined dispatcher
+        // 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 (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
+        }
+        // Otherwise, it was some other dispatcher that successfully dispatched the coroutine
     }
-    // It was some other dispatcher that successfully dispatched the coroutine
     COROUTINE_SUSPENDED
 }
 

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

@@ -71,9 +71,10 @@ public class RecoverableTestCancellationException(message: String? = null) : Can
 public fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
     val dispatcher = context[ContinuationInterceptor] as CoroutineDispatcher
     return object : CoroutineDispatcher() {
-        override fun dispatch(context: CoroutineContext, block: Runnable) {
+        override fun isDispatchNeeded(context: CoroutineContext): Boolean =
+            dispatcher.isDispatchNeeded(context)
+        override fun dispatch(context: CoroutineContext, block: Runnable) =
             dispatcher.dispatch(context, block)
-        }
     }
 }