Fix possible leaking of CompletedExceptionally on racy cancellation

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

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

@@ -80,19 +80,20 @@ private inline fun <T> startDirect(completion: Continuation<T>, block: () -> Any
  */
 internal fun <T, R> AbstractCoroutine<T>.startUndispatchedOrReturn(receiver: R, block: suspend R.() -> T): Any? {
     initParentJob()
-    return undispatchedResult({ true }) { block.startCoroutineUninterceptedOrReturn(receiver, this) }
+    return undispatchedResult({ true }) {
+        block.startCoroutineUninterceptedOrReturn(receiver, this)
+    }
 }
 
 /**
  * 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? {
+    receiver: R, block: suspend R.() -> T): Any? {
     initParentJob()
-    return undispatchedResult({ e -> !(e is TimeoutCancellationException && e.coroutine === timeoutCoroutine) })
-    { block.startCoroutineUninterceptedOrReturn(receiver, this) }
+    return undispatchedResult({ e -> !(e is TimeoutCancellationException && e.coroutine === this) }) {
+        block.startCoroutineUninterceptedOrReturn(receiver, this)
+    }
 }
 
 private inline fun <T> AbstractCoroutine<T>.undispatchedResult(
@@ -124,10 +125,10 @@ private inline fun <T> AbstractCoroutine<T>.undispatchedResult(
         makeCompletingOnce(result, MODE_IGNORE) -> {
             val state = state
             if (state is CompletedExceptionally) {
-                if (shouldThrow(state.cause)) {
-                    throw state.cause
-                } else {
-                    result
+                when {
+                    shouldThrow(state.cause) -> throw state.cause
+                    result is CompletedExceptionally -> throw result.cause
+                    else -> result
                 }
             } else {
                 state

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

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

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

@@ -24,7 +24,7 @@ class WithTimeoutOrNullJvmTest : TestBase() {
 
     @Test
     fun testIgnoredTimeout() = runTest {
-        val value =withTimeout(1) {
+        val value = withTimeout(1) {
             Thread.sleep(10)
             42
         }
@@ -41,4 +41,27 @@ class WithTimeoutOrNullJvmTest : TestBase() {
 
         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)
+    }
 }