Do not report exceptions raised in CoroutineDispatcher.dispatch as internal errors

Author: zuevmaxim

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

@@ -191,7 +191,7 @@ internal class DispatchedContinuation<in T>(
         if (dispatcher.isDispatchNeeded(context)) {
             _state = state
             resumeMode = MODE_ATOMIC
-            dispatcher.dispatch(context, this)
+            dispatchWithExceptionHandling(context)
         } else {
             executeUnconfined(state, MODE_ATOMIC) {
                 withCoroutineContext(this.context, countOrElement) {
@@ -212,7 +212,7 @@ internal class DispatchedContinuation<in T>(
         if (dispatcher.isDispatchNeeded(context)) {
             _state = state
             resumeMode = MODE_CANCELLABLE
-            dispatcher.dispatch(context, this)
+            dispatchWithExceptionHandling(context)
         } else {
             executeUnconfined(state, MODE_CANCELLABLE) {
                 if (!resumeCancelled(state)) {
@@ -260,6 +260,14 @@ 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)
+        }
+    }
 }
 
 /**

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

@@ -78,6 +78,7 @@ internal abstract class DispatchedTask<in T> internal constructor(
         assert { resumeMode != MODE_UNINITIALIZED } // should have been set before dispatching
         val taskContext = this.taskContext
         var fatalException: Throwable? = null
+        var dispatchException: DispatchException? = null
         try {
             val delegate = delegate as DispatchedContinuation<T>
             val continuation = delegate.continuation
@@ -103,12 +104,17 @@ internal abstract class DispatchedTask<in T> internal constructor(
                     }
                 }
             }
+        } catch (e: DispatchException) {
+            dispatchException = e
         } catch (e: Throwable) {
             // This instead of runCatching to have nicer stacktrace and debug experience
             fatalException = e
         } finally {
             val result = runCatching { taskContext.afterTask() }
             handleFatalException(fatalException, result.exceptionOrNull())
+            if (dispatchException != null) {
+                handleCoroutineException(delegate.context, dispatchException.cause!!)
+            }
         }
     }
 
@@ -213,3 +219,10 @@ internal inline fun DispatchedTask<*>.runUnconfinedEventLoop(
 internal inline fun Continuation<*>.resumeWithStackTrace(exception: Throwable) {
     resumeWith(Result.failure(recoverStackTrace(exception, this)))
 }
+
+/**
+ * This exception holds an exception raised in [CoroutineDispatcher.dispatch] method
+ *
+ * @see DispatchedContinuation.dispatchWithExceptionHandling
+ */
+internal class DispatchException(cause: Throwable) : Exception(cause)

kotlinx-coroutines-core/jvm/test/ExecutorsTest.kt

@@ -119,4 +119,79 @@ class ExecutorsTest : TestBase() {
         dispatcher.close()
         check(executorService.isShutdown)
     }
+
+    @Test
+    fun testEarlyExecutorShutdown() {
+        runTestExceptionInDispatch(6, { it is RejectedExecutionException }) {
+            expect(1)
+            val dispatcher = newSingleThreadContext("Ctx")
+            launch(dispatcher) {
+                withContext(Dispatchers.Default) {
+                    expect(2)
+                    delay(100)
+                    expect(4)
+                }
+            }
+
+            delay(50)
+            expect(3)
+
+            dispatcher.close()
+        }
+    }
+
+    @Test
+    fun testExceptionInDispatch() {
+        runTestExceptionInDispatch(5, { it is TestException }) {
+            val dispatcher = object : CoroutineDispatcher() {
+                private var closed = false
+                override fun dispatch(context: CoroutineContext, block: Runnable) {
+                    if (closed) throw TestException()
+                    Dispatchers.Default.dispatch(context, block)
+                }
+
+                fun close() {
+                    closed = true
+                }
+            }
+            launch(dispatcher) {
+                withContext(Dispatchers.Default) {
+                    expect(1)
+                    delay(100)
+                    expect(3)
+                }
+            }
+
+            delay(50)
+            expect(2)
+            dispatcher.close()
+        }
+    }
+
+    private fun runTestExceptionInDispatch(
+        totalSteps: Int,
+        isExpectedException: (Throwable) -> Boolean,
+        block: suspend CoroutineScope.() -> Unit,
+    ) {
+        var mainThread: Thread? = null
+        val exceptionHandler = CoroutineExceptionHandler { _, e ->
+            if (isExpectedException(e)) {
+                expect(totalSteps - 1)
+                mainThread!!.run {
+                    interrupt()
+                    unpark(this)
+                }
+            } else {
+                expectUnreached()
+            }
+        }
+        try {
+            runBlocking(exceptionHandler) {
+                block()
+                mainThread = Thread.currentThread()
+            }
+        } catch (_: InterruptedException) {
+            finish(totalSteps)
+        }
+    }
 }