From 2d0686b9818c834a1344737afd0805afd2367845 Mon Sep 17 00:00:00 2001
From: Louis CAD <louis.cognault@gmail.com>
Date: Wed, 2 Sep 2020 02:04:28 +0200
Subject: [PATCH] Add awaitCancellation (#2225)

Fixes #2213
---
 .../api/kotlinx-coroutines-core.api           |  1 +
 kotlinx-coroutines-core/common/src/Delay.kt   | 41 +++++++++++++++++++
 .../common/test/AwaitCancellationTest.kt      | 27 ++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 kotlinx-coroutines-core/common/test/AwaitCancellationTest.kt

diff --git a/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api b/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
index 36cbdb6960..5df13700d0 100644
--- a/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
+++ b/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
@@ -267,6 +267,7 @@ public final class kotlinx/coroutines/Delay$DefaultImpls {
 }
 
 public final class kotlinx/coroutines/DelayKt {
+	public static final fun awaitCancellation (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun delay (JLkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun delay-p9JZ4hM (DLkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
diff --git a/kotlinx-coroutines-core/common/src/Delay.kt b/kotlinx-coroutines-core/common/src/Delay.kt
index ab80912269..ff0a6fc73b 100644
--- a/kotlinx-coroutines-core/common/src/Delay.kt
+++ b/kotlinx-coroutines-core/common/src/Delay.kt
@@ -58,12 +58,51 @@ public interface Delay {
         DefaultDelay.invokeOnTimeout(timeMillis, block)
 }
 
+/**
+ * Suspends until cancellation, in which case it will throw a [CancellationException].
+ *
+ * This function returns [Nothing], so it can be used in any coroutine,
+ * regardless of the required return type.
+ *
+ * Usage example in callback adapting code:
+ *
+ * ```kotlin
+ * fun currentTemperature(): Flow<Temperature> = callbackFlow {
+ *     val callback = SensorCallback { degreesCelsius: Double ->
+ *         trySend(Temperature.celsius(degreesCelsius))
+ *     }
+ *     try {
+ *         registerSensorCallback(callback)
+ *         awaitCancellation() // Suspends to keep getting updates until cancellation.
+ *     } finally {
+ *         unregisterSensorCallback(callback)
+ *     }
+ * }
+ * ```
+ *
+ * Usage example in (non declarative) UI code:
+ *
+ * ```kotlin
+ * suspend fun showStuffUntilCancelled(content: Stuff): Nothing {
+ *     someSubView.text = content.title
+ *     anotherSubView.text = content.description
+ *     someView.visibleInScope {
+ *         awaitCancellation() // Suspends so the view stays visible.
+ *     }
+ * }
+ * ```
+ */
+@ExperimentalCoroutinesApi
+public suspend fun awaitCancellation(): Nothing = suspendCancellableCoroutine {}
+
 /**
  * Delays coroutine for a given time without blocking a thread and resumes it after a specified time.
  * This suspending function is cancellable.
  * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException].
  *
+ * If you want to delay forever (until cancellation), consider using [awaitCancellation] instead.
+ *
  * Note that delay can be used in [select] invocation with [onTimeout][SelectBuilder.onTimeout] clause.
  *
  * Implementation note: how exactly time is tracked is an implementation detail of [CoroutineDispatcher] in the context.
@@ -82,6 +121,8 @@ public suspend fun delay(timeMillis: Long) {
  * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException].
  *
+ * If you want to delay forever (until cancellation), consider using [awaitCancellation] instead.
+ *
  * Note that delay can be used in [select] invocation with [onTimeout][SelectBuilder.onTimeout] clause.
  *
  * Implementation note: how exactly time is tracked is an implementation detail of [CoroutineDispatcher] in the context.
diff --git a/kotlinx-coroutines-core/common/test/AwaitCancellationTest.kt b/kotlinx-coroutines-core/common/test/AwaitCancellationTest.kt
new file mode 100644
index 0000000000..2fe0c914e6
--- /dev/null
+++ b/kotlinx-coroutines-core/common/test/AwaitCancellationTest.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlin.test.*
+
+class AwaitCancellationTest : TestBase() {
+
+    @Test
+    fun testCancellation() = runTest(expected = { it is CancellationException }) {
+        expect(1)
+        coroutineScope {
+            val deferred: Deferred<Nothing> = async {
+                expect(2)
+                awaitCancellation()
+            }
+            yield()
+            expect(3)
+            require(deferred.isActive)
+            deferred.cancel()
+            finish(4)
+            deferred.await()
+        }
+    }
+}