Merge pull request #791 from Kotlin/undispatched-cancellation

Use final Job state in undispatched coroutine completion

Author: qwwdfsad

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

@@ -74,7 +74,7 @@ private fun <U, T: U> setupTimeout(
     coroutine.disposeOnCompletion(context.delay.invokeOnTimeout(coroutine.time, coroutine))
     // restart block using new coroutine with new job,
     // however start it as undispatched coroutine, because we are already in the proper context
-    return coroutine.startUndispatchedOrReturn(coroutine, block)
+    return coroutine.startUndispatchedOrReturnIgnoreTimeout(coroutine, block)
 }
 
 private open class TimeoutCoroutine<U, in T: U>(

common/kotlinx-coroutines-core-common/src/intrinsics/Undispatched.kt

@@ -70,29 +70,70 @@ private inline fun <T> startDirect(completion: Continuation<T>, block: () -> Any
 }
 
 /**
- * Starts this coroutine with the given code [block] in the same context and returns result when it
+ * Starts this coroutine with the given code [block] in the same context and returns coroutine result when it
  * completes without suspension.
  * This function shall be invoked at most once on this coroutine.
+ * This function checks cancellation of the outer [Job] on fast-path.
  *
  * First, this function initializes parent job from the `parentContext` of this coroutine that was passed to it
  * during construction. Second, it starts the coroutine using [startCoroutineUninterceptedOrReturn].
  */
 internal fun <T, R> AbstractCoroutine<T>.startUndispatchedOrReturn(receiver: R, block: suspend R.() -> T): Any? {
     initParentJob()
-    return undispatchedResult { block.startCoroutineUninterceptedOrReturn(receiver, this) }
+    return undispatchedResult({ true }) {
+        block.startCoroutineUninterceptedOrReturn(receiver, this)
+    }
 }
 
-private inline fun <T> AbstractCoroutine<T>.undispatchedResult(startBlock: () -> Any?): Any? {
+/**
+ * Same as [startUndispatchedOrReturn], but ignores [TimeoutCancellationException] on fast-path.
+ */
+internal fun <T, R> AbstractCoroutine<T>.startUndispatchedOrReturnIgnoreTimeout(
+    receiver: R, block: suspend R.() -> T): Any? {
+    initParentJob()
+    return undispatchedResult({ e -> !(e is TimeoutCancellationException && e.coroutine === this) }) {
+        block.startCoroutineUninterceptedOrReturn(receiver, this)
+    }
+}
+
+private inline fun <T> AbstractCoroutine<T>.undispatchedResult(
+    shouldThrow: (Throwable) -> Boolean,
+    startBlock: () -> Any?
+): Any? {
     val result = try {
         startBlock()
     } catch (e: Throwable) {
         CompletedExceptionally(e)
     }
+
+    /*
+     * We're trying to complete our undispatched block here and have three code-paths:
+     * 1) Suspended.
+     *
+     * Or we are completing our block (and its job).
+     * 2) If we can't complete it, we suspend, probably waiting for children (2)
+     * 3) If we have successfully completed the whole coroutine here in an undispatched manner,
+     *    we should decide which result to return. We have two options: either return proposed update or actual final state.
+     *    But if fact returning proposed value is not an option, otherwise we will ignore possible cancellation or child failure.
+     *
+     * shouldThrow parameter is a special code path for timeout coroutine:
+     * If timeout is exceeded, but withTimeout() block was not suspended, we would like to return block value,
+     * not a timeout exception.
+     */
     return when {
         result === COROUTINE_SUSPENDED -> COROUTINE_SUSPENDED
         makeCompletingOnce(result, MODE_IGNORE) -> {
-            if (result is CompletedExceptionally) throw result.cause else result
+            val state = state
+            if (state is CompletedExceptionally) {
+                when {
+                    shouldThrow(state.cause) -> throw state.cause
+                    result is CompletedExceptionally -> throw result.cause
+                    else -> result
+                }
+            } else {
+                state
+            }
         }
         else -> COROUTINE_SUSPENDED
     }
-}
+}

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

@@ -4,6 +4,8 @@
 
 package kotlinx.coroutines
 
+import kotlin.coroutines.*
+
 public expect open class TestBase constructor() {
     public val isStressTest: Boolean
     public val stressTestMultiplier: Int
@@ -26,3 +28,13 @@ public class TestException1(message: String? = null) : Throwable(message)
 public class TestException2(message: String? = null) : Throwable(message)
 public class TestException3(message: String? = null) : Throwable(message)
 public class TestRuntimeException(message: String? = null) : RuntimeException(message)
+
+// Wrap context to avoid fast-paths on dispatcher comparison
+public fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
+    val dispatcher = context[ContinuationInterceptor] as CoroutineDispatcher
+    return object : CoroutineDispatcher() {
+        override fun dispatch(context: CoroutineContext, block: Runnable) {
+            dispatcher.dispatch(context, block)
+        }
+    }
+}

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class UndispatchedResultTest : TestBase() {
+
+    @Test
+    fun testWithContext() = runTest {
+        invokeTest { block -> withContext(wrapperDispatcher(coroutineContext), block) }
+    }
+
+    @Test
+    fun testWithContextFastPath() = runTest {
+        invokeTest { block -> withContext(coroutineContext, block) }
+    }
+
+    @Test
+    fun testWithTimeout() = runTest {
+        invokeTest { block -> withTimeout(Long.MAX_VALUE, block) }
+    }
+
+    @Test
+    fun testAsync() = runTest {
+        invokeTest { block -> async(NonCancellable, block = block).await() }
+    }
+
+    @Test
+    fun testCoroutineScope() = runTest {
+        invokeTest { block -> coroutineScope(block) }
+    }
+
+    private suspend fun invokeTest(scopeProvider: suspend (suspend CoroutineScope.() -> Unit) -> Unit) {
+        invokeTest(EmptyCoroutineContext, scopeProvider)
+        invokeTest(Unconfined, scopeProvider)
+    }
+
+    private suspend fun invokeTest(
+        context: CoroutineContext,
+        scopeProvider: suspend (suspend CoroutineScope.() -> Unit) -> Unit
+    ) {
+        try {
+            scopeProvider { block(context) }
+        } catch (e: TestException) {
+            finish(5)
+            reset()
+        }
+    }
+
+    private suspend fun CoroutineScope.block(context: CoroutineContext) {
+        try {
+            expect(1)
+            // Will cancel its parent
+            async(context) {
+                expect(2)
+                throw TestException()
+            }.await()
+        } catch (e: TestException) {
+            expect(3)
+        }
+        expect(4)
+    }
+}

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

@@ -7,7 +7,6 @@
 
 package kotlinx.coroutines
 
-import kotlin.coroutines.*
 import kotlin.test.*
 
 class WithContextTest : TestBase() {
@@ -85,20 +84,25 @@ class WithContextTest : TestBase() {
         }
         expect(2)
         val job = Job()
-        val result = withContext(coroutineContext + job) { // same context + new job
-            expect(3) // still here
-            job.cancel() // cancel out job!
-            try {
-                yield() // shall throw CancellationException
-                expectUnreached()
-            } catch (e: CancellationException) {
-                expect(4)
+        try {
+            withContext(coroutineContext + job) {
+                // same context + new job
+                expect(3) // still here
+                job.cancel() // cancel out job!
+                try {
+                    yield() // shall throw CancellationException
+                    expectUnreached()
+                } catch (e: CancellationException) {
+                    expect(4)
+                }
+                "OK"
             }
-            "OK"
+
+            expectUnreached()
+        } catch (e: CancellationException) {
+            expect(5)
+            // will wait for the first coroutine
         }
-        assertEquals("OK", result)
-        expect(5)
-        // will wait for the first coroutine
     }
 
     @Test
@@ -300,13 +304,4 @@ class WithContextTest : TestBase() {
         }
         finish(5)
     }
-
-    private fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
-        val dispatcher = context[ContinuationInterceptor] as CoroutineDispatcher
-        return object : CoroutineDispatcher() {
-            override fun dispatch(context: CoroutineContext, block: Runnable) {
-                dispatcher.dispatch(context, block)
-            }
-        }
-    }
 }

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

@@ -238,4 +238,4 @@ class WithTimeoutOrNullTest : TestBase() {
             finish(4)
         }
     }
-}
+}

core/kotlinx-coroutines-core/test/WithTimeoutOrNullJvmTest.kt

@@ -21,4 +21,47 @@ class WithTimeoutOrNullJvmTest : TestBase() {
         // outer timeout results in null
         assertEquals(null, result)
     }
-}
+
+    @Test
+    fun testIgnoredTimeout() = runTest {
+        val value = withTimeout(1) {
+            Thread.sleep(10)
+            42
+        }
+
+        assertEquals(42, value)
+    }
+
+    @Test
+    fun testIgnoredTimeoutOnNull() = runTest {
+        val value = withTimeoutOrNull(1) {
+            Thread.sleep(10)
+            42
+        }
+
+        assertEquals(42, value)
+    }
+
+    @Test
+    fun testIgnoredTimeoutOnNullThrowsCancellation() = runTest {
+        try {
+            withTimeoutOrNull(1) {
+                expect(1)
+                Thread.sleep(10)
+                throw CancellationException()
+            }
+            expectUnreached()
+        } catch (e: CancellationException) {
+            finish(2)
+        }
+    }
+
+    @Test
+    fun testIgnoredTimeoutOnNullThrowsOnYield() = runTest {
+        val value = withTimeoutOrNull(1) {
+            Thread.sleep(10)
+            yield()
+        }
+        assertNull(value)
+    }
+}