Implement withTimeoutOrNull via withTimeout to avoid timing bugs and races.

Remove deprecated API

Fixes #498

Author: qwwdfsad

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

@@ -69,13 +69,6 @@ private fun <U, T: U> setupTimeout(
     return coroutine.startUndispatchedOrReturn(coroutine, block)
 }
 
-/**
- * @suppress **Deprecated**: for binary compatibility only
- */
-@Deprecated("for binary compatibility only", level=DeprecationLevel.HIDDEN)
-public suspend fun <T> withTimeout(time: Long, unit: TimeUnit = TimeUnit.MILLISECONDS, block: suspend () -> T): T =
-    withTimeout(time, unit) { block() }
-
 private open class TimeoutCoroutine<U, in T: U>(
     @JvmField val time: Long,
     @JvmField val unit: TimeUnit,
@@ -140,32 +133,21 @@ public suspend fun <T> withTimeoutOrNull(time: Int, block: suspend CoroutineScop
  */
 public suspend fun <T> withTimeoutOrNull(time: Long, unit: TimeUnit = TimeUnit.MILLISECONDS, block: suspend CoroutineScope.() -> T): T? {
     if (time <= 0L) return null
-    return suspendCoroutineUninterceptedOrReturn { uCont ->
-        setupTimeout(TimeoutOrNullCoroutine(time, unit, uCont), block)
-    }
-}
-
-/**
- * @suppress **Deprecated**: for binary compatibility only
- */
-@Deprecated("for binary compatibility only", level=DeprecationLevel.HIDDEN)
-public suspend fun <T> withTimeoutOrNull(time: Long, unit: TimeUnit = TimeUnit.MILLISECONDS, block: suspend () -> T): T? =
-    withTimeoutOrNull(time, unit) { block() }
 
-private class TimeoutOrNullCoroutine<T>(
-    time: Long,
-    unit: TimeUnit,
-    uCont: Continuation<T?> // unintercepted continuation
-) : TimeoutCoroutine<T?, T>(time, unit, uCont) {
-    @Suppress("UNCHECKED_CAST")
-    internal override fun onCompletionInternal(state: Any?, mode: Int) {
-        if (state is CompletedExceptionally) {
-            val exception = state.cause
-            if (exception is TimeoutCancellationException && exception.coroutine === this)
-                uCont.resumeUninterceptedMode(null, mode) else
-                uCont.resumeUninterceptedWithExceptionMode(exception, mode)
-        } else
-            uCont.resumeUninterceptedMode(state as T, mode)
+    var coroutine: TimeoutCoroutine<T?, T?>? = null
+    try {
+        return suspendCoroutineUninterceptedOrReturn { uCont ->
+            TimeoutCoroutine(time, unit, uCont).let {
+                coroutine = it
+                setupTimeout<T?, T?>(it, block)
+            }
+        }
+    } catch (e: TimeoutCancellationException) {
+        // Return null iff it's our exception, otherwise propagate it upstream (e.g. in case of nested withTimeouts)
+        if (e.coroutine === coroutine) {
+            return null
+        }
+        throw e
     }
 }
 

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

@@ -7,10 +7,12 @@
 
 package kotlinx.coroutines.experimental
 
+import kotlinx.coroutines.experimental.channels.*
 import kotlin.coroutines.experimental.*
 import kotlin.test.*
 
 class WithTimeoutOrNullTest : TestBase() {
+
     /**
      * Tests a case of no timeout and no suspension inside.
      */
@@ -81,6 +83,23 @@ class WithTimeoutOrNullTest : TestBase() {
         finish(2)
     }
 
+    @Test
+    fun testSmallTimeout() = runTest {
+        val channel = Channel<Int>(1)
+        val value = withTimeoutOrNull(1) {
+            channel.receive()
+        }
+
+        assertNull(value)
+    }
+
+    @Test
+    fun testThrowException() = runTest(expected = {it is AssertionError}) {
+        withTimeoutOrNull(Long.MAX_VALUE) {
+            throw AssertionError()
+        }
+    }
+
     @Test
     fun testInnerTimeoutTest() = runTest(
         expected = { it is CancellationException }
@@ -96,6 +115,19 @@ class WithTimeoutOrNullTest : TestBase() {
         expectUnreached() // will timeout
     }
 
+    @Test
+    fun testNestedTimeout() = runTest(expected = { it is TimeoutCancellationException }) {
+        withTimeoutOrNull(Long.MAX_VALUE) {
+            // Exception from this withTimeout is not suppressed by withTimeoutOrNull
+            withTimeout(10) {
+                delay(Long.MAX_VALUE)
+                1
+            }
+        }
+
+        expectUnreached()
+    }
+
     @Test
     fun testOuterTimeoutTest() = runTest {
         var counter = 0