From 089e163ca805e4280afb2d4837cf483e0ad02383 Mon Sep 17 00:00:00 2001
From: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>
Date: Thu, 27 May 2021 17:48:03 +0300
Subject: [PATCH] Fail eagerly during exceptions in isDispatchNeeded

That helps to pro-actively catch cases like #2717 and to report such exception in even more robust manner
---
 .../common/src/intrinsics/Cancellable.kt          | 15 ++++++++++++++-
 .../jvm/test/FailFastOnStartTest.kt               | 12 +++++++++++-
 .../jvm/test/FailingCoroutinesMachineryTest.kt    | 14 ++++++++++----
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/kotlinx-coroutines-core/common/src/intrinsics/Cancellable.kt b/kotlinx-coroutines-core/common/src/intrinsics/Cancellable.kt
index 173f0afb65..f5b96a8d88 100644
--- a/kotlinx-coroutines-core/common/src/intrinsics/Cancellable.kt
+++ b/kotlinx-coroutines-core/common/src/intrinsics/Cancellable.kt
@@ -49,6 +49,19 @@ private inline fun runSafely(completion: Continuation<*>, block: () -> Unit) {
     try {
         block()
     } catch (e: Throwable) {
-        completion.resumeWith(Result.failure(e))
+        dispatcherFailure(completion, e)
     }
 }
+
+private fun dispatcherFailure(completion: Continuation<*>, e: Throwable) {
+    /*
+     * This method is invoked when we failed to start a coroutine due to the throwing
+     * dispatcher implementation or missing Dispatchers.Main.
+     * This situation is not recoverable, so we are trying to deliver the exception by all means:
+     * 1) Resume the coroutine with an exception, so it won't prevent its parent from completion
+     * 2) Rethrow the exception immediately, so it will crash the caller (e.g. when the coroutine had
+     *    no parent or it was async/produce over MainScope).
+     */
+    completion.resumeWith(Result.failure(e))
+    throw e
+}
diff --git a/kotlinx-coroutines-core/jvm/test/FailFastOnStartTest.kt b/kotlinx-coroutines-core/jvm/test/FailFastOnStartTest.kt
index 15cb83ceed..8a7878c9a6 100644
--- a/kotlinx-coroutines-core/jvm/test/FailFastOnStartTest.kt
+++ b/kotlinx-coroutines-core/jvm/test/FailFastOnStartTest.kt
@@ -70,8 +70,18 @@ class FailFastOnStartTest : TestBase() {
         val actor = actor<Int>(Dispatchers.Main, start = CoroutineStart.LAZY) { fail() }
         actor.send(1)
     }
-    
+
     private fun mainException(e: Throwable): Boolean {
         return e is IllegalStateException && e.message?.contains("Module with the Main dispatcher is missing") ?: false
     }
+
+    @Test
+    fun testProduceNonChild() = runTest(expected = ::mainException) {
+        produce<Int>(Job() + Dispatchers.Main) { fail() }
+    }
+
+    @Test
+    fun testAsyncNonChild() = runTest(expected = ::mainException) {
+        async<Int>(Job() + Dispatchers.Main) { fail() }
+    }
 }
diff --git a/kotlinx-coroutines-core/jvm/test/FailingCoroutinesMachineryTest.kt b/kotlinx-coroutines-core/jvm/test/FailingCoroutinesMachineryTest.kt
index c9f722a5b8..04b0ba547d 100644
--- a/kotlinx-coroutines-core/jvm/test/FailingCoroutinesMachineryTest.kt
+++ b/kotlinx-coroutines-core/jvm/test/FailingCoroutinesMachineryTest.kt
@@ -33,7 +33,7 @@ class FailingCoroutinesMachineryTest(
 
     private var caught: Throwable? = null
     private val latch = CountDownLatch(1)
-    private var exceptionHandler = CoroutineExceptionHandler { _, t -> caught = t;latch.countDown() }
+    private var exceptionHandler = CoroutineExceptionHandler { _, t -> caught = t; latch.countDown() }
     private val lazyOuterDispatcher = lazy { newFixedThreadPoolContext(1, "") }
 
     private object FailingUpdate : ThreadContextElement<Unit> {
@@ -115,14 +115,20 @@ class FailingCoroutinesMachineryTest(
 
     @Test
     fun testElement() = runTest {
-        launch(NonCancellable + dispatcher.value + exceptionHandler + element) {}
+        // Top-level throwing dispatcher may rethrow an exception right here
+        runCatching {
+            launch(NonCancellable + dispatcher.value + exceptionHandler + element) {}
+        }
         checkException()
     }
 
     @Test
     fun testNestedElement() = runTest {
-        launch(NonCancellable + dispatcher.value + exceptionHandler) {
-            launch(element) {  }
+        // Top-level throwing dispatcher may rethrow an exception right here
+        runCatching {
+            launch(NonCancellable + dispatcher.value + exceptionHandler) {
+                launch(element) { }
+            }
         }
         checkException()
     }