Use final Job state in undispatched coroutine completion

  * When using in-place coroutine builders, startUndispatchedOrReturn will return block's final value on fast-path. This is not correct because if any of the children had failed, whether or not exception is thrown depend only on child timings. Moreover, in the face of such race exception may be ignored if scope owner wasn't cancelled or is missing.
  * Introduce special hack for withTimeout and withTimeoutOrNull to ignore TimeoutCancellationException on the fast path
  * withContext, coroutineScope, supervisorScope, withTimeout and withTimeoutOrNull are affected and now check cancellation on fast-path

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, coroutine)
 }
 
 private open class TimeoutCoroutine<U, in T: U>(

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

@@ -70,29 +70,69 @@ 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,
+    timeoutCoroutine: Job
+): Any? {
+    initParentJob()
+    return undispatchedResult({ e -> !(e is TimeoutCancellationException && e.coroutine === timeoutCoroutine) })
+    { 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) {
+                if (shouldThrow(state.cause)) {
+                    throw state.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)
-            }
-        }
-    }
 }

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

@@ -21,4 +21,24 @@ 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)
+    }
+}