Protect isDispatchNeeded calls

Author: zuevmaxim

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

@@ -192,7 +192,7 @@ public abstract class CoroutineDispatcher :
      * @suppress **This an internal API and should not be used from general code.**
      */
     @InternalCoroutinesApi
-    public open fun dispatchYield(context: CoroutineContext, block: Runnable): Unit = dispatch(context, block)
+    public open fun dispatchYield(context: CoroutineContext, block: Runnable): Unit = safeDispatch(context, block)
 
     /**
      * Returns a continuation that wraps the provided [continuation], thus intercepting all resumptions.

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

@@ -26,7 +26,7 @@ public suspend fun yield(): Unit = suspendCoroutineUninterceptedOrReturn sc@ { u
     val context = uCont.context
     context.ensureActive()
     val cont = uCont.intercepted() as? DispatchedContinuation<Unit> ?: return@sc Unit
-    if (cont.dispatcher.isDispatchNeeded(context)) {
+    if (cont.dispatcher.safeIsDispatchNeeded(context)) {
         // this is a regular dispatcher -- do simple dispatchYield
         cont.dispatchYield(context, Unit)
     } else {

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

@@ -188,10 +188,10 @@ internal class DispatchedContinuation<in T>(
     override fun resumeWith(result: Result<T>) {
         val context = continuation.context
         val state = result.toState()
-        if (dispatcher.isDispatchNeeded(context)) {
+        if (dispatcher.safeIsDispatchNeeded(context)) {
             _state = state
             resumeMode = MODE_ATOMIC
-            dispatchWithExceptionHandling(context)
+            dispatcher.safeDispatch(context, this)
         } else {
             executeUnconfined(state, MODE_ATOMIC) {
                 withCoroutineContext(this.context, countOrElement) {
@@ -209,10 +209,10 @@ internal class DispatchedContinuation<in T>(
         noinline onCancellation: ((cause: Throwable) -> Unit)?
     ) {
         val state = result.toState(onCancellation)
-        if (dispatcher.isDispatchNeeded(context)) {
+        if (dispatcher.safeIsDispatchNeeded(context)) {
             _state = state
             resumeMode = MODE_CANCELLABLE
-            dispatchWithExceptionHandling(context)
+            dispatcher.safeDispatch(context, this)
         } else {
             executeUnconfined(state, MODE_CANCELLABLE) {
                 if (!resumeCancelled(state)) {
@@ -260,13 +260,21 @@ internal class DispatchedContinuation<in T>(
 
     override fun toString(): String =
         "DispatchedContinuation[$dispatcher, ${continuation.toDebugString()}]"
+}
 
-    private fun dispatchWithExceptionHandling(context: CoroutineContext) {
-        try {
-            dispatcher.dispatch(context, this)
-        } catch (e: Throwable) {
-            throw DispatchException(e, dispatcher, context)
-        }
+internal fun CoroutineDispatcher.safeDispatch(context: CoroutineContext, runnable: Runnable) {
+    try {
+        dispatch(context, runnable)
+    } catch (e: Throwable) {
+        throw DispatchException(e, this, context)
+    }
+}
+
+internal fun CoroutineDispatcher.safeIsDispatchNeeded(context: CoroutineContext): Boolean {
+    try {
+        return isDispatchNeeded(context)
+    } catch (e: Throwable) {
+        throw DispatchException(e, this, context)
     }
 }
 

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

@@ -157,8 +157,8 @@ internal fun <T> DispatchedTask<T>.dispatch(mode: Int) {
         // dispatch directly using this instance's Runnable implementation
         val dispatcher = delegate.dispatcher
         val context = delegate.context
-        if (dispatcher.isDispatchNeeded(context)) {
-            dispatcher.dispatch(context, this)
+        if (dispatcher.safeIsDispatchNeeded(context)) {
+            dispatcher.safeDispatch(context, this)
         } else {
             resumeUnconfined()
         }
@@ -223,7 +223,7 @@ internal inline fun Continuation<*>.resumeWithStackTrace(exception: Throwable) {
 /**
  * This exception holds an exception raised in [CoroutineDispatcher.dispatch] method
  *
- * @see DispatchedContinuation.dispatchWithExceptionHandling
+ * @see safeDispatch
  */
 internal class DispatchException(
     override val cause: Throwable,

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

@@ -42,7 +42,7 @@ internal class LimitedDispatcher(
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
         dispatchInternal(block) { worker ->
-            dispatcher.dispatch(this, worker)
+            dispatcher.safeDispatch(this, worker)
         }
     }
 
@@ -116,10 +116,10 @@ internal class LimitedDispatcher(
                 }
                 currentTask = obtainTaskOrDeallocateWorker() ?: return
                 // 16 is our out-of-thin-air constant to emulate fairness. Used in JS dispatchers as well
-                if (++fairnessCounter >= 16 && dispatcher.isDispatchNeeded(this@LimitedDispatcher)) {
+                if (++fairnessCounter >= 16 && dispatcher.safeIsDispatchNeeded(this@LimitedDispatcher)) {
                     // Do "yield" to let other views execute their runnable as well
                     // Note that we do not decrement 'runningWorkers' as we are still committed to our part of work
-                    dispatcher.dispatch(this@LimitedDispatcher, this)
+                    dispatcher.safeDispatch(this@LimitedDispatcher, this)
                     return
                 }
             }

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

@@ -12,9 +12,9 @@ internal class NamedDispatcher(
     private val name: String
 ) : CoroutineDispatcher(), Delay by (dispatcher as? Delay ?: DefaultDelay) {
 
-    override fun isDispatchNeeded(context: CoroutineContext): Boolean = dispatcher.isDispatchNeeded(context)
+    override fun isDispatchNeeded(context: CoroutineContext): Boolean = dispatcher.safeIsDispatchNeeded(context)
 
-    override fun dispatch(context: CoroutineContext, block: Runnable) = dispatcher.dispatch(context, block)
+    override fun dispatch(context: CoroutineContext, block: Runnable) = dispatcher.safeDispatch(context, block)
 
     @InternalCoroutinesApi
     override fun dispatchYield(context: CoroutineContext, block: Runnable) = dispatcher.dispatchYield(context, block)

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

@@ -106,8 +106,8 @@ public fun CoroutineDispatcher.asExecutor(): Executor =
 
 private class DispatcherExecutor(@JvmField val dispatcher: CoroutineDispatcher) : Executor {
     override fun execute(block: Runnable) {
-        if (dispatcher.isDispatchNeeded(EmptyCoroutineContext)) {
-            dispatcher.dispatch(EmptyCoroutineContext, block)
+        if (dispatcher.safeIsDispatchNeeded(EmptyCoroutineContext)) {
+            dispatcher.safeDispatch(EmptyCoroutineContext, block)
         } else {
             block.run()
         }