Make withContext cancellable on return (instead of atomically cancell…

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

@@ -156,7 +156,7 @@ public suspend fun <T> withContext(
         }
     }
     // SLOW PATH -- use new dispatcher
-    val coroutine = DispatchedCoroutine(newContext, uCont) // MODE_ATOMIC_DEFAULT
+    val coroutine = DispatchedCoroutine(newContext, uCont) // MODE_CANCELLABLE
     coroutine.initParentJob()
     block.startCoroutineCancellable(coroutine, coroutine)
     coroutine.getResult()
@@ -215,7 +215,7 @@ private class DispatchedCoroutine<in T>(
     context: CoroutineContext,
     uCont: Continuation<T>
 ) : ScopeCoroutine<T>(context, uCont) {
-    override val defaultResumeMode: Int get() = MODE_ATOMIC_DEFAULT
+    override val defaultResumeMode: Int get() = MODE_CANCELLABLE
 
     // this is copy-and-paste of a decision state machine inside AbstractionContinuation
     // todo: we may some-how abstract it via inline class

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

@@ -218,32 +218,35 @@ internal abstract class DispatchedTask<in T>(
 
     public final override fun run() {
         val taskContext = this.taskContext
-        var exception: Throwable? = null
+        var fatalException: Throwable? = null
         try {
             val delegate = delegate as DispatchedContinuation<T>
             val continuation = delegate.continuation
             val context = continuation.context
-            val job = if (resumeMode.isCancellableMode) context[Job] else null
             val state = takeState() // NOTE: Must take state in any case, even if cancelled
             withCoroutineContext(context, delegate.countOrElement) {
-                if (job != null && !job.isActive) {
+                val exception = getExceptionalResult(state)
+                val job = if (resumeMode.isCancellableMode) context[Job] else null
+                /*
+                 * Check whether continuation was originally resumed with an exception.
+                 * If so, it dominates cancellation, otherwise the original exception
+                 * will be silently lost.
+                 */
+                if (exception == null && job != null && !job.isActive) {
                     val cause = job.getCancellationException()
                     cancelResult(state, cause)
                     continuation.resumeWithStackTrace(cause)
                 } else {
-                    val exception = getExceptionalResult(state)
-                    if (exception != null)
-                        continuation.resumeWithStackTrace(exception)
-                    else
-                        continuation.resume(getSuccessfulResult(state))
+                    if (exception != null) continuation.resumeWithStackTrace(exception)
+                    else continuation.resume(getSuccessfulResult(state))
                 }
             }
         } catch (e: Throwable) {
             // This instead of runCatching to have nicer stacktrace and debug experience
-            exception = e
+            fatalException = e
         } finally {
             val result = runCatching { taskContext.afterTask() }
-            handleFatalException(exception, result.exceptionOrNull())
+            handleFatalException(fatalException, result.exceptionOrNull())
         }
     }
 

kotlinx-coroutines-core/common/test/TestBase.common.kt

@@ -74,3 +74,5 @@ public fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
     }
 }
 
+public suspend fun wrapperDispatcher(): CoroutineContext = wrapperDispatcher(coroutineContext)
+

kotlinx-coroutines-core/common/test/WithContextTest.kt

@@ -15,7 +15,7 @@ class WithContextTest : TestBase() {
     fun testThrowException() = runTest {
         expect(1)
         try {
-            withContext(coroutineContext) {
+            withContext<Unit>(coroutineContext) {
                 expect(2)
                 throw AssertionError()
             }
@@ -31,7 +31,7 @@ class WithContextTest : TestBase() {
     fun testThrowExceptionFromWrappedContext() = runTest {
         expect(1)
         try {
-            withContext(wrapperDispatcher(coroutineContext)) {
+            withContext<Unit>(wrapperDispatcher(coroutineContext)) {
                 expect(2)
                 throw AssertionError()
             }
@@ -151,7 +151,7 @@ class WithContextTest : TestBase() {
             expect(2)
             try {
                 // Same dispatcher, different context
-                withContext(CoroutineName("testRunCancellationUndispatchedVsException")) {
+                withContext<Unit>(CoroutineName("testRunCancellationUndispatchedVsException")) {
                     expect(3)
                     yield() // must suspend
                     expect(5)
@@ -176,7 +176,7 @@ class WithContextTest : TestBase() {
             expect(2)
             try {
                 // "Different" dispatcher (schedules execution back and forth)
-                withContext(wrapperDispatcher(coroutineContext)) {
+                withContext<Unit>(wrapperDispatcher(coroutineContext)) {
                     expect(4)
                     yield() // must suspend
                     expect(6)
@@ -204,7 +204,7 @@ class WithContextTest : TestBase() {
         job = launch(Job()) {
             try {
                 expect(3)
-                withContext(wrapperDispatcher(coroutineContext)) {
+                withContext<Unit>(wrapperDispatcher(coroutineContext)) {
                     require(isActive)
                     expect(5)
                     job!!.cancel()
@@ -349,6 +349,26 @@ class WithContextTest : TestBase() {
         expectUnreached()
     }
 
+    @Test
+    fun testSequentialCancellation() = runTest {
+        val job = launch {
+            expect(1)
+            withContext(wrapperDispatcher()) {
+                expect(2)
+            }
+            expectUnreached()
+        }
+
+        yield()
+        val job2 = launch {
+            expect(3)
+            job.cancel()
+        }
+
+        joinAll(job, job2)
+        finish(4)
+    }
+
     private class Wrapper(val value: String) : Incomplete {
         override val isActive: Boolean
             get() =  error("")