Introduce CancellableContinuation.resume with onCancelling lambda

* Allows safe return of closeable resources from suspending
functions, as it provides a way to close a resource if the
corresponding job was cancelled.
* Documentation on the context and expected behavior of
CompletionHandler implementations is updated.

Fixes #1044

Author: elizarov

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -40,6 +40,7 @@ public abstract interface class kotlinx/coroutines/CancellableContinuation : kot
 	public abstract fun isActive ()Z
 	public abstract fun isCancelled ()Z
 	public abstract fun isCompleted ()Z
+	public abstract fun resume (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;)V
 	public abstract fun resumeUndispatched (Lkotlinx/coroutines/CoroutineDispatcher;Ljava/lang/Object;)V
 	public abstract fun resumeUndispatchedWithException (Lkotlinx/coroutines/CoroutineDispatcher;Ljava/lang/Throwable;)V
 	public abstract fun tryResume (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;
@@ -54,6 +55,7 @@ public final class kotlinx/coroutines/CancellableContinuation$DefaultImpls {
 public class kotlinx/coroutines/CancellableContinuationImpl : kotlin/coroutines/jvm/internal/CoroutineStackFrame, kotlinx/coroutines/CancellableContinuation {
 	public fun <init> (Lkotlin/coroutines/Continuation;I)V
 	public fun cancel (Ljava/lang/Throwable;)Z
+	public fun cancelResult (Ljava/lang/Object;Ljava/lang/Throwable;)V
 	public fun completeResume (Ljava/lang/Object;)V
 	public fun getCallerFrame ()Lkotlin/coroutines/jvm/internal/CoroutineStackFrame;
 	public fun getContext ()Lkotlin/coroutines/CoroutineContext;
@@ -68,6 +70,7 @@ public class kotlinx/coroutines/CancellableContinuationImpl : kotlin/coroutines/
 	public fun isCancelled ()Z
 	public fun isCompleted ()Z
 	protected fun nameString ()Ljava/lang/String;
+	public fun resume (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;)V
 	public fun resumeUndispatched (Lkotlinx/coroutines/CoroutineDispatcher;Ljava/lang/Object;)V
 	public fun resumeUndispatchedWithException (Lkotlinx/coroutines/CoroutineDispatcher;Ljava/lang/Throwable;)V
 	public fun resumeWith (Ljava/lang/Object;)V

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

@@ -123,10 +123,15 @@ public interface CancellableContinuation<in T> : Continuation<T> {
      * with cancellation exception. Otherwise, the handler will be invoked once on cancellation if this
      * continuation is cancelled.
      *
-     * Installed [handler] should not throw any exceptions. If it does, they will get caught,
-     * wrapped into [CompletionHandlerException], and rethrown, potentially causing the crash of unrelated code.
+     * Installed [handler] should not throw any exceptions.
+     * If it does, they will get caught, wrapped into [CompletionHandlerException] and
+     * processed as uncaught exception in the context of the current coroutine
+     * (see [CoroutineExceptionHandler]).
      *
      * At most one [handler] can be installed on one continuation.
+     *
+     * **Note**: Implementation of `CompletionHandler` must be fast, non-blocking, and thread-safe.
+     * There is no guarantee on the execution context in which the [handler] is invoked.
      */
     public fun invokeOnCancellation(handler: CompletionHandler)
 
@@ -151,6 +156,33 @@ public interface CancellableContinuation<in T> : Continuation<T> {
      */
     @ExperimentalCoroutinesApi
     public fun CoroutineDispatcher.resumeUndispatchedWithException(exception: Throwable)
+
+    /**
+     * Resumes this continuation with a given [value] and calls the specified [onCancellation]
+     * handler when resumed too late (when continuation was already cancelled) or when resumed
+     * successfully (before cancellation), but coroutine's job was cancelled before it had a
+     * chance to run in its dispatcher, so that suspended function threw an exception
+     * instead of returning this value.
+     *
+     * Installed [onCancellation] handler should not throw any exceptions.
+     * If it does, they will get caught, wrapped into [CompletionHandlerException] and
+     * processed as uncaught exception in the context of the current coroutine
+     * (see [CoroutineExceptionHandler]).
+     *
+     * This function shall be used when resuming with a resource that must be closed by the
+     * code that had called the corresponding suspending function, e.g.:
+     *
+     * ```
+     * continuation.resume(resource) {
+     *     resource.close()
+     * }
+     * ```
+     *
+     * **Note**: Implementation of `CompletionHandler` must be fast, non-blocking, and thread-safe.
+     * There is no guarantee on the execution context in which the [onCancellation] handler is invoked.
+     */
+    @ExperimentalCoroutinesApi // since 1.2.0, tentatively graduates in 1.3.0
+    public fun resume(value: T, onCancellation: CompletionHandler)
 }
 
 /**

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

@@ -102,6 +102,14 @@ internal open class CancellableContinuationImpl<in T>(
 
     override fun takeState(): Any? = state
 
+    override fun cancelResult(state: Any?, cause: Throwable) {
+        if (state is CompletedWithCancellation) {
+            invokeHandlerSafely {
+                state.onCancellation.invokeIt(cause)
+            }
+        }
+    }
+
     public override fun cancel(cause: Throwable?): Boolean {
         _state.loop { state ->
             if (state !is NotCompleted) return false // false if already complete or cancelling
@@ -165,8 +173,18 @@ internal open class CancellableContinuationImpl<in T>(
         return getSuccessfulResult(state)
     }
 
-    override fun resumeWith(result: Result<T>) =
+    override fun resumeWith(result: Result<T>) {
         resumeImpl(result.toState(), resumeMode)
+    }
+
+    override fun resume(value: T, onCancellation: CompletionHandler) {
+        if (!resumeImpl(CompletedWithCancellation(value, onCancellation), resumeMode)) {
+            // too late to resume (was cancelled) -- call handler
+            invokeHandlerSafely {
+                onCancellation.invokeIt((state as? CompletedExceptionally)?.cause)
+            }
+        }
+    }
 
     internal fun resumeWithExceptionMode(exception: Throwable, mode: Int) =
         resumeImpl(CompletedExceptionally(exception), mode)
@@ -219,22 +237,23 @@ internal open class CancellableContinuationImpl<in T>(
         dispatch(mode)
     }
 
-    private fun resumeImpl(proposedUpdate: Any?, resumeMode: Int) {
+    // returns true when successfully dispatched resumed, false if too late
+    private fun resumeImpl(proposedUpdate: Any?, resumeMode: Int): Boolean {
         _state.loop { state ->
             when (state) {
                 is NotCompleted -> {
                     if (!_state.compareAndSet(state, proposedUpdate)) return@loop // retry on cas failure
                     disposeParentHandle()
                     dispatchResume(resumeMode)
-                    return
+                    return true
                 }
                 is CancelledContinuation -> {
                     /*
                      * If continuation was cancelled, then resume attempt must be ignored,
                      * because cancellation is asynchronous and may race with resume.
                      * Racy exceptions will be lost, too.
                      */
-                    if (state.makeResumed()) return // ok -- resumed just once
+                    if (state.makeResumed()) return false // ok -- tried to resume just once
                 }
             }
             alreadyResumedError(proposedUpdate) // otherwise -- an error (second resume attempt)
@@ -307,7 +326,11 @@ internal open class CancellableContinuationImpl<in T>(
 
     @Suppress("UNCHECKED_CAST")
     override fun <T> getSuccessfulResult(state: Any?): T =
-        if (state is CompletedIdempotentResult) state.result as T else state as T
+        when (state) {
+            is CompletedIdempotentResult -> state.result as T
+            is CompletedWithCancellation -> state.result as T
+            else -> state as T
+        }
 
     // For nicer debugging
     public override fun toString(): String =
@@ -344,3 +367,11 @@ private class CompletedIdempotentResult(
 ) {
     override fun toString(): String = "CompletedIdempotentResult[$result]"
 }
+
+private class CompletedWithCancellation(
+    @JvmField val result: Any?,
+    @JvmField val onCancellation: CompletionHandler
+) {
+    override fun toString(): String = "CompletedWithCancellation[$result]"
+}
+

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

@@ -207,6 +207,8 @@ internal abstract class DispatchedTask<in T>(
 
     public abstract fun takeState(): Any?
 
+    public open fun cancelResult(state: Any?, cause: Throwable) {}
+
     @Suppress("UNCHECKED_CAST")
     public open fun <T> getSuccessfulResult(state: Any?): T =
         state as T
@@ -224,9 +226,11 @@ internal abstract class DispatchedTask<in T>(
             val job = if (resumeMode.isCancellableMode) context[Job] else null
             val state = takeState() // NOTE: Must take state in any case, even if cancelled
             withCoroutineContext(context, delegate.countOrElement) {
-                if (job != null && !job.isActive)
-                    continuation.resumeWithException(job.getCancellationException())
-                else {
+                if (job != null && !job.isActive) {
+                    val cause = job.getCancellationException()
+                    cancelResult(state, cause)
+                    continuation.resumeWithException(cause)
+                } else {
                     val exception = getExceptionalResult(state)
                     if (exception != null)
                         continuation.resumeWithStackTrace(exception)

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

@@ -273,7 +273,8 @@ public interface Job : CoroutineContext.Element {
      * Installed [handler] should not throw any exceptions. If it does, they will get caught,
      * wrapped into [CompletionHandlerException], and rethrown, potentially causing crash of unrelated code.
      *
-     * **Note**: Implementations of `CompletionHandler` must be fast and _lock-free_.
+     * **Note**: Implementation of `CompletionHandler` must be fast, non-blocking, and thread-safe.
+     * There is no guarantee on the execution context in which the [handler] is invoked.
      */
     public fun invokeOnCompletion(handler: CompletionHandler): DisposableHandle
 
@@ -304,7 +305,8 @@ public interface Job : CoroutineContext.Element {
      * **Note**: This function is a part of internal machinery that supports parent-child hierarchies
      * and allows for implementation of suspending functions that wait on the Job's state.
      * This function should not be used in general application code.
-     * Implementations of `CompletionHandler` must be fast and _lock-free_.
+     * Implementation of `CompletionHandler` must be fast, non-blocking, and thread-safe.
+     * There is no guarantee on the execution context in which the [handler] is invoked.
      *
      * @param onCancelling when `true`, then the [handler] is invoked as soon as this job transitions to _cancelling_ state;
      *        when `false` then the [handler] is invoked only when it transitions to _completed_ state.

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

@@ -0,0 +1,122 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+@file:Suppress("NAMED_ARGUMENTS_NOT_ALLOWED") // KT-21913
+
+package kotlinx.coroutines
+
+import kotlin.test.*
+
+/**
+ * Test for [CancellableContinuation.resume] with `onCancellation` parameter.
+ */
+class CancellableResumeTest : TestBase() {
+    @Test
+    fun testResumeImmediateNormally() = runTest {
+        expect(1)
+        val ok = suspendCancellableCoroutine<String> { cont ->
+            expect(2)
+            cont.invokeOnCancellation { expectUnreached() }
+            cont.resume("OK") { expectUnreached() }
+            expect(3)
+        }
+        assertEquals("OK", ok)
+        finish(4)
+    }
+
+    @Test
+    fun testResumeImmediateAfterCancel() = runTest(
+        expected = { it is TestException }
+    ) {
+        expect(1)
+        val ok = suspendCancellableCoroutine<String> { cont ->
+            expect(2)
+            cont.invokeOnCancellation { expect(3) }
+            cont.cancel(TestException("FAIL"))
+            expect(4)
+            cont.resume("OK") { cause ->
+                expect(5)
+                assertTrue(cause is TestException)
+            }
+            finish(6)
+        }
+        expectUnreached()
+    }
+
+    @Test
+    fun testResumeLaterNormally() = runTest {
+        expect(1)
+        lateinit var cc: CancellableContinuation<String>
+        launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(2)
+            val ok = suspendCancellableCoroutine<String> { cont ->
+                expect(3)
+                cont.invokeOnCancellation { expectUnreached() }
+                cc = cont
+            }
+            assertEquals("OK", ok)
+            finish(6)
+        }
+        expect(4)
+        cc.resume("OK") { expectUnreached() }
+        expect(5)
+    }
+
+    @Test
+    fun testResumeLaterAfterCancel() = runTest {
+        expect(1)
+        lateinit var cc: CancellableContinuation<String>
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(2)
+            try {
+                suspendCancellableCoroutine<String> { cont ->
+                    expect(3)
+                    cont.invokeOnCancellation { expect(5) }
+                    cc = cont
+                }
+                expectUnreached()
+            } catch (e: CancellationException) {
+                finish(9)
+            }
+        }
+        expect(4)
+        job.cancel(TestCancellationException())
+        expect(6)
+        cc.resume("OK") { cause ->
+            expect(7)
+            assertTrue(cause is TestCancellationException)
+        }
+        expect(8)
+    }
+
+    @Test
+    fun testResumeCancelWhileDispatched() = runTest {
+        expect(1)
+        lateinit var cc: CancellableContinuation<String>
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(2)
+            try {
+                suspendCancellableCoroutine<String> { cont ->
+                    expect(3)
+                    // resumed first, then cancelled, so no invokeOnCancellation call
+                    cont.invokeOnCancellation { expectUnreached() }
+                    cc = cont
+                }
+                expectUnreached()
+            } catch (e: CancellationException) {
+                expect(8)
+            }
+        }
+        expect(4)
+        cc.resume("OK") { cause ->
+            expect(7)
+            assertTrue(cause is TestCancellationException)
+        }
+        expect(5)
+        job.cancel(TestCancellationException()) // cancel while execution is dispatched
+        expect(6)
+        yield() // to coroutine -- throws cancellation exception
+        finish(9)
+    }
+}