Implement withTimeoutOrNull via withTimeout to avoid timing bugs and …

…races. Remove deprecated API Fixes #498
Kotlin · Aug 21, 2018 · 988eb26 · 988eb26
1 parent 44f52b8
commit 988eb26
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 36 deletions.
diff --git a/binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt b/binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt
@@ -416,14 +416,10 @@ public final class kotlinx/coroutines/experimental/RunnableKt {
 
 public final class kotlinx/coroutines/experimental/ScheduledKt {
 	public static final fun withTimeout (ILkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
-	public static final synthetic fun withTimeout (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
 	public static final fun withTimeout (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
-	public static synthetic fun withTimeout$default (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/experimental/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static synthetic fun withTimeout$default (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static final fun withTimeoutOrNull (ILkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
-	public static final synthetic fun withTimeoutOrNull (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
 	public static final fun withTimeoutOrNull (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;)Ljava/lang/Object;
-	public static synthetic fun withTimeoutOrNull$default (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/experimental/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static synthetic fun withTimeoutOrNull$default (JLjava/util/concurrent/TimeUnit;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/experimental/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 }
 

diff --git a/common/kotlinx-coroutines-core-common/src/Scheduled.kt b/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
     }
 }
 

diff --git a/common/kotlinx-coroutines-core-common/test/WithTimeoutOrNullTest.kt b/common/kotlinx-coroutines-core-common/test/WithTimeoutOrNullTest.kt
@@ -7,6 +7,7 @@
 
 package kotlinx.coroutines.experimental
 
+import kotlinx.coroutines.experimental.channels.*
 import kotlin.coroutines.experimental.*
 import kotlin.test.*
 
@@ -81,6 +82,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 +114,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