Recover stacktraces for no-dispatched continuations, so recovery works in 'suspend fun main' cases to further improve user experience

Fixes #1328

Author: qwwdfsad

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

@@ -307,7 +307,14 @@ internal fun <T> DispatchedTask<T>.resume(delegate: Continuation<T>, useMode: In
     val state = takeState()
     val exception = getExceptionalResult(state)
     if (exception != null) {
-        delegate.resumeWithExceptionMode(exception, useMode)
+        /*
+         * Recover stacktrace for non-dispatched tasks.
+         * We usually do not recover stacktrace in a `resume` as all resumes go through `DispatchedTask.run`
+         * and we recover stacktraces there, but this is not the case for a `suspend fun main()` that knows nothing about
+         * kotlinx.coroutines and DispatchedTask
+         */
+        val recovered = if (delegate is DispatchedTask<*>) exception else recoverStackTrace(exception, delegate)
+        delegate.resumeWithExceptionMode(recovered, useMode)
     } else {
         delegate.resumeMode(getSuccessfulResult(state), useMode)
     }

kotlinx-coroutines-core/jvm/test/exceptions/StackTraceRecoveryTest.kt

@@ -6,9 +6,11 @@ package kotlinx.coroutines.exceptions
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.selects.*
 import org.junit.Test
 import java.util.concurrent.*
+import kotlin.coroutines.*
 import kotlin.test.*
 
 /*
@@ -271,4 +273,40 @@ class StackTraceRecoveryTest : TestBase() {
             checkCycles(e)
         }
     }
+
+
+    private suspend fun throws() {
+        yield() // TCE
+        throw RecoverableTestException()
+    }
+
+    private suspend fun awaiter() {
+        val task = GlobalScope.async(Dispatchers.Default, start = CoroutineStart.LAZY) { throws() }
+        task.await()
+        yield() // TCE
+    }
+
+    @Test
+    fun testNonDispatchedRecovery() {
+        val await = suspend { awaiter() }
+
+        val barrier = CyclicBarrier(2)
+        var exception: Throwable? = null
+        await.startCoroutineUnintercepted(Continuation(EmptyCoroutineContext) {
+            exception = it.exceptionOrNull()
+            barrier.await()
+        })
+
+        barrier.await()
+        val e = exception
+        assertNotNull(e)
+        verifyStackTrace(e, "kotlinx.coroutines.RecoverableTestException\n" +
+                "\tat kotlinx.coroutines.exceptions.StackTraceRecoveryTest.throws(StackTraceRecoveryTest.kt:280)\n" +
+                "\tat kotlinx.coroutines.exceptions.StackTraceRecoveryTest\$throws\$1.invokeSuspend(StackTraceRecoveryTest.kt)\n" +
+                "\t(Coroutine boundary)\n" +
+                "\tat kotlinx.coroutines.DeferredCoroutine.await\$suspendImpl(Builders.common.kt:99)\n" +
+                "\tat kotlinx.coroutines.exceptions.StackTraceRecoveryTest.awaiter(StackTraceRecoveryTest.kt:285)\n" +
+                "\tat kotlinx.coroutines.exceptions.StackTraceRecoveryTest\$testNonDispatchedRecovery\$await\$1.invokeSuspend(StackTraceRecoveryTest.kt:291)\n" +
+                "Caused by: kotlinx.coroutines.RecoverableTestException")
+    }
 }