Add await for task created with cancellation token

integration/kotlinx-coroutines-play-services/README.md

@@ -6,6 +6,7 @@ Extension functions:
 
 | **Name** | **Description**
 | -------- | ---------------
+| [((CancellationToken) -> Task).await][((CancellationToken) -> Task).await] | Awaits for completion of a Task constructed with a CancellationToken
 | [Task.await][await] | Awaits for completion of the Task (cancellable)
 | [Deferred.asTask][asTask] | Converts a deferred value to a Task
 
@@ -25,5 +26,6 @@ val snapshot = try {
 // Do stuff
 ```
 
+[((CancellationToken) -> Task).await]: https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-play-services/kotlinx.coroutines.tasks/awaitConstructedTask.html
 [await]: https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-play-services/kotlinx.coroutines.tasks/com.google.android.gms.tasks.-task/await.html
 [asTask]: https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-play-services/kotlinx.coroutines.tasks/kotlinx.coroutines.-deferred/as-task.html

integration/kotlinx-coroutines-play-services/api/kotlinx-coroutines-play-services.api

@@ -2,5 +2,6 @@ public final class kotlinx/coroutines/tasks/TasksKt {
 	public static final fun asDeferred (Lcom/google/android/gms/tasks/Task;)Lkotlinx/coroutines/Deferred;
 	public static final fun asTask (Lkotlinx/coroutines/Deferred;)Lcom/google/android/gms/tasks/Task;
 	public static final fun await (Lcom/google/android/gms/tasks/Task;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun await (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 

integration/kotlinx-coroutines-play-services/src/Tasks.kt

@@ -6,15 +6,9 @@
 
 package kotlinx.coroutines.tasks
 
-import com.google.android.gms.tasks.CancellationTokenSource
-import com.google.android.gms.tasks.RuntimeExecutionException
-import com.google.android.gms.tasks.Task
-import com.google.android.gms.tasks.TaskCompletionSource
-import kotlinx.coroutines.CancellationException
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.Deferred
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.suspendCancellableCoroutine
+import com.google.android.gms.tasks.*
+import kotlinx.coroutines.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 
 /**
@@ -105,3 +99,31 @@ public suspend fun <T> Task<T>.await(): T {
         }
     }
 }
+
+/**
+ * Awaits for completion of the created task without blocking a thread.
+ *
+ * Prefer this method over [Task.await] if a [Task] can be constructed with a [CancellationToken], to cancel the task
+ * if this function is cancelled.
+ *
+ * This suspending function is cancellable.
+ * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
+ * stops waiting for the completion stage and immediately resumes with [CancellationException].
+ *
+ * @see [Task.await]
+ */
+@OptIn(ExperimentalContracts::class)
+@ExperimentalCoroutinesApi
+public suspend fun <T> (suspend (CancellationToken) -> Task<T>).await(): T {
+    contract { callsInPlace(this@await, InvocationKind.EXACTLY_ONCE) }
+
+    val cancellation = CancellationTokenSource()
+    val task = this(cancellation.token)
+
+    return try {
+        task.await()
+    } catch (cancellationException: CancellationException) {
+        cancellation.cancel()
+        throw cancellationException
+    }
+}

integration/kotlinx-coroutines-play-services/test/TaskTest.kt

@@ -149,5 +149,69 @@ class TaskTest : TestBase() {
         }
     }
 
+    @Test
+    fun testAwaitConstructedTask() = runTest {
+        var taskCompletionSource: TaskCompletionSource<Int>? = null
+
+        val deferred: Deferred<Int> = async(start = CoroutineStart.UNDISPATCHED) {
+            val taskCreator: suspend (CancellationToken) -> Task<Int> = { cancellationToken ->
+                taskCompletionSource = TaskCompletionSource<Int>(cancellationToken)
+                taskCompletionSource!!.task
+            }
+            taskCreator.await()
+        }
+
+        assertFalse(deferred.isCompleted)
+        taskCompletionSource!!.setResult(42)
+
+        assertEquals(42, deferred.await())
+        assertTrue(deferred.isCompleted)
+    }
+
+    @Test
+    fun testFailedAwaitConstructedTask() {
+        var taskCompletionSource: TaskCompletionSource<Int>? = null
+
+        val deferred: Deferred<Int> = GlobalScope.async(start = CoroutineStart.UNDISPATCHED) {
+            val taskCreator: suspend (CancellationToken) -> Task<Int> = { cancellationToken ->
+                taskCompletionSource = TaskCompletionSource<Int>(cancellationToken)
+                taskCompletionSource!!.task
+            }
+            taskCreator.await()
+        }
+
+        assertFalse(deferred.isCompleted)
+        taskCompletionSource!!.setException(TestException("something went wrong"))
+
+        runTest(expected = { it is TestException }) {
+            deferred.await()
+        }
+    }
+
+    @Test
+    fun testCancelledAwaitConstructedTask() = runTest {
+        var task: Task<Int>? = null
+
+        val deferred: Deferred<Int> = async(start = CoroutineStart.UNDISPATCHED) {
+            val taskCreator: suspend (CancellationToken) -> Task<Int> = { cancellationToken ->
+                task = TaskCompletionSource<Int>(cancellationToken).task
+                task!!
+            }
+            taskCreator.await()
+        }
+
+        assertFalse(deferred.isCompleted)
+        deferred.cancel()
+
+        try {
+            deferred.await()
+            fail("deferred.await() should be cancelled")
+        } catch (e: Exception) {
+            assertTrue(e is CancellationException)
+        }
+
+        assertTrue(task!!.isCanceled)
+    }
+
     class TestException(message: String) : Exception(message)
 }