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.

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

@@ -20,23 +20,30 @@ import kotlin.coroutines.intrinsics.*
  *
  * If the coroutine dispatcher is [Unconfined][Dispatchers.Unconfined], this
  * functions suspends only when there are other unconfined coroutines working and forming an event-loop.
- * For other dispatchers, this function does not call [CoroutineDispatcher.isDispatchNeeded] and
- * always suspends to be resumed later. If there is no [CoroutineDispatcher] in the context, it does not suspend.
+ * For other dispatchers, this function calls [CoroutineDispatcher.dispatch] and
+ * always suspends to be resumed later regardless of the result of [CoroutineDispatcher.isDispatchNeeded].
+ * If there is no [CoroutineDispatcher] in the context, it does not suspend.
  */
 public suspend fun yield(): Unit = suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
  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)
- }
  }
 }