Stabilize CancellableContinuation.resume with onCancellation (#4090)

Additionally, give `onCancellation` extra parameters
that help to avoid allocating closures in some cases.

Fixes #4088

Author: dkhalanskyjb

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -39,10 +39,11 @@ public abstract interface class kotlinx/coroutines/CancellableContinuation : kot
 	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 resume (Ljava/lang/Object;Lkotlin/jvm/functions/Function3;)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;
-	public abstract fun tryResume (Ljava/lang/Object;Ljava/lang/Object;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
+	public abstract fun tryResume (Ljava/lang/Object;Ljava/lang/Object;Lkotlin/jvm/functions/Function3;)Ljava/lang/Object;
 	public abstract fun tryResumeWithException (Ljava/lang/Throwable;)Ljava/lang/Object;
 }
 
@@ -54,7 +55,7 @@ public final class kotlinx/coroutines/CancellableContinuation$DefaultImpls {
 public class kotlinx/coroutines/CancellableContinuationImpl : kotlin/coroutines/jvm/internal/CoroutineStackFrame, kotlinx/coroutines/CancellableContinuation, kotlinx/coroutines/Waiter {
 	public fun <init> (Lkotlin/coroutines/Continuation;I)V
 	public final fun callCancelHandler (Lkotlinx/coroutines/CancelHandler;Ljava/lang/Throwable;)V
-	public final fun callOnCancellation (Lkotlin/jvm/functions/Function1;Ljava/lang/Throwable;)V
+	public final fun callOnCancellation (Lkotlin/jvm/functions/Function3;Ljava/lang/Throwable;Ljava/lang/Object;)V
 	public fun cancel (Ljava/lang/Throwable;)Z
 	public fun completeResume (Ljava/lang/Object;)V
 	public fun getCallerFrame ()Lkotlin/coroutines/jvm/internal/CoroutineStackFrame;
@@ -70,12 +71,13 @@ public class kotlinx/coroutines/CancellableContinuationImpl : kotlin/coroutines/
 	public fun isCompleted ()Z
 	protected fun nameString ()Ljava/lang/String;
 	public fun resume (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;)V
+	public fun resume (Ljava/lang/Object;Lkotlin/jvm/functions/Function3;)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
 	public fun toString ()Ljava/lang/String;
 	public fun tryResume (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;
-	public fun tryResume (Ljava/lang/Object;Ljava/lang/Object;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
+	public fun tryResume (Ljava/lang/Object;Ljava/lang/Object;Lkotlin/jvm/functions/Function3;)Ljava/lang/Object;
 	public fun tryResumeWithException (Ljava/lang/Throwable;)Ljava/lang/Object;
 }
 

kotlinx-coroutines-core/api/kotlinx-coroutines-core.klib.api

@@ -104,13 +104,14 @@ abstract interface <#A: in kotlin/Any?> kotlinx.coroutines.channels/SendChannel
 abstract interface <#A: in kotlin/Any?> kotlinx.coroutines/CancellableContinuation : kotlin.coroutines/Continuation<#A> { // kotlinx.coroutines/CancellableContinuation|null[0]
     abstract fun (kotlinx.coroutines/CoroutineDispatcher).resumeUndispatched(#A) // kotlinx.coroutines/CancellableContinuation.resumeUndispatched|[email protected](1:0){}[0]
     abstract fun (kotlinx.coroutines/CoroutineDispatcher).resumeUndispatchedWithException(kotlin/Throwable) // kotlinx.coroutines/CancellableContinuation.resumeUndispatchedWithException|resumeUndispatchedWithException@kotlinx.coroutines.CoroutineDispatcher(kotlin.Throwable){}[0]
+    abstract fun <#A1: #A> resume(#A1, kotlin/Function3<kotlin/Throwable, #A1, kotlin.coroutines/CoroutineContext, kotlin/Unit>?) // kotlinx.coroutines/CancellableContinuation.resume|resume(0:0;kotlin.Function3<kotlin.Throwable,0:0,kotlin.coroutines.CoroutineContext,kotlin.Unit>?){0§<1:0>}[0]
+    abstract fun <#A1: #A> tryResume(#A1, kotlin/Any?, kotlin/Function3<kotlin/Throwable, #A1, kotlin.coroutines/CoroutineContext, kotlin/Unit>?): kotlin/Any? // kotlinx.coroutines/CancellableContinuation.tryResume|tryResume(0:0;kotlin.Any?;kotlin.Function3<kotlin.Throwable,0:0,kotlin.coroutines.CoroutineContext,kotlin.Unit>?){0§<1:0>}[0]
     abstract fun cancel(kotlin/Throwable? = ...): kotlin/Boolean // kotlinx.coroutines/CancellableContinuation.cancel|cancel(kotlin.Throwable?){}[0]
     abstract fun completeResume(kotlin/Any) // kotlinx.coroutines/CancellableContinuation.completeResume|completeResume(kotlin.Any){}[0]
     abstract fun initCancellability() // kotlinx.coroutines/CancellableContinuation.initCancellability|initCancellability(){}[0]
     abstract fun invokeOnCancellation(kotlin/Function1<kotlin/Throwable?, kotlin/Unit>) // kotlinx.coroutines/CancellableContinuation.invokeOnCancellation|invokeOnCancellation(kotlin.Function1<kotlin.Throwable?,kotlin.Unit>){}[0]
     abstract fun resume(#A, kotlin/Function1<kotlin/Throwable, kotlin/Unit>?) // kotlinx.coroutines/CancellableContinuation.resume|resume(1:0;kotlin.Function1<kotlin.Throwable,kotlin.Unit>?){}[0]
     abstract fun tryResume(#A, kotlin/Any? = ...): kotlin/Any? // kotlinx.coroutines/CancellableContinuation.tryResume|tryResume(1:0;kotlin.Any?){}[0]
-    abstract fun tryResume(#A, kotlin/Any?, kotlin/Function1<kotlin/Throwable, kotlin/Unit>?): kotlin/Any? // kotlinx.coroutines/CancellableContinuation.tryResume|tryResume(1:0;kotlin.Any?;kotlin.Function1<kotlin.Throwable,kotlin.Unit>?){}[0]
     abstract fun tryResumeWithException(kotlin/Throwable): kotlin/Any? // kotlinx.coroutines/CancellableContinuation.tryResumeWithException|tryResumeWithException(kotlin.Throwable){}[0]
     abstract val isActive // kotlinx.coroutines/CancellableContinuation.isActive|{}isActive[0]
         abstract fun <get-isActive>(): kotlin/Boolean // kotlinx.coroutines/CancellableContinuation.isActive.<get-isActive>|<get-isActive>(){}[0]
@@ -774,11 +775,13 @@ open annotation class kotlinx.coroutines/ObsoleteCoroutinesApi : kotlin/Annotati
 }
 open class <#A: in kotlin/Any?> kotlinx.coroutines/CancellableContinuationImpl : kotlinx.coroutines.internal/CoroutineStackFrame, kotlinx.coroutines/CancellableContinuation<#A>, kotlinx.coroutines/DispatchedTask<#A>, kotlinx.coroutines/Waiter { // kotlinx.coroutines/CancellableContinuationImpl|null[0]
     constructor <init>(kotlin.coroutines/Continuation<#A>, kotlin/Int) // kotlinx.coroutines/CancellableContinuationImpl.<init>|<init>(kotlin.coroutines.Continuation<1:0>;kotlin.Int){}[0]
+    final fun <#A1: kotlin/Any?> callOnCancellation(kotlin/Function3<kotlin/Throwable, #A1, kotlin.coroutines/CoroutineContext, kotlin/Unit>, kotlin/Throwable, #A1) // kotlinx.coroutines/CancellableContinuationImpl.callOnCancellation|callOnCancellation(kotlin.Function3<kotlin.Throwable,0:0,kotlin.coroutines.CoroutineContext,kotlin.Unit>;kotlin.Throwable;0:0){0§<kotlin.Any?>}[0]
     final fun callCancelHandler(kotlinx.coroutines/CancelHandler, kotlin/Throwable?) // kotlinx.coroutines/CancellableContinuationImpl.callCancelHandler|callCancelHandler(kotlinx.coroutines.CancelHandler;kotlin.Throwable?){}[0]
-    final fun callOnCancellation(kotlin/Function1<kotlin/Throwable, kotlin/Unit>, kotlin/Throwable) // kotlinx.coroutines/CancellableContinuationImpl.callOnCancellation|callOnCancellation(kotlin.Function1<kotlin.Throwable,kotlin.Unit>;kotlin.Throwable){}[0]
     final fun getResult(): kotlin/Any? // kotlinx.coroutines/CancellableContinuationImpl.getResult|getResult(){}[0]
     open fun (kotlinx.coroutines/CoroutineDispatcher).resumeUndispatched(#A) // kotlinx.coroutines/CancellableContinuationImpl.resumeUndispatched|[email protected](1:0){}[0]
     open fun (kotlinx.coroutines/CoroutineDispatcher).resumeUndispatchedWithException(kotlin/Throwable) // kotlinx.coroutines/CancellableContinuationImpl.resumeUndispatchedWithException|resumeUndispatchedWithException@kotlinx.coroutines.CoroutineDispatcher(kotlin.Throwable){}[0]
+    open fun <#A1: #A> resume(#A1, kotlin/Function3<kotlin/Throwable, #A1, kotlin.coroutines/CoroutineContext, kotlin/Unit>?) // kotlinx.coroutines/CancellableContinuationImpl.resume|resume(0:0;kotlin.Function3<kotlin.Throwable,0:0,kotlin.coroutines.CoroutineContext,kotlin.Unit>?){0§<1:0>}[0]
+    open fun <#A1: #A> tryResume(#A1, kotlin/Any?, kotlin/Function3<kotlin/Throwable, #A1, kotlin.coroutines/CoroutineContext, kotlin/Unit>?): kotlin/Any? // kotlinx.coroutines/CancellableContinuationImpl.tryResume|tryResume(0:0;kotlin.Any?;kotlin.Function3<kotlin.Throwable,0:0,kotlin.coroutines.CoroutineContext,kotlin.Unit>?){0§<1:0>}[0]
     open fun cancel(kotlin/Throwable?): kotlin/Boolean // kotlinx.coroutines/CancellableContinuationImpl.cancel|cancel(kotlin.Throwable?){}[0]
     open fun completeResume(kotlin/Any) // kotlinx.coroutines/CancellableContinuationImpl.completeResume|completeResume(kotlin.Any){}[0]
     open fun getContinuationCancellationCause(kotlinx.coroutines/Job): kotlin/Throwable // kotlinx.coroutines/CancellableContinuationImpl.getContinuationCancellationCause|getContinuationCancellationCause(kotlinx.coroutines.Job){}[0]
@@ -791,7 +794,6 @@ open class <#A: in kotlin/Any?> kotlinx.coroutines/CancellableContinuationImpl :
     open fun resumeWith(kotlin/Result<#A>) // kotlinx.coroutines/CancellableContinuationImpl.resumeWith|resumeWith(kotlin.Result<1:0>){}[0]
     open fun toString(): kotlin/String // kotlinx.coroutines/CancellableContinuationImpl.toString|toString(){}[0]
     open fun tryResume(#A, kotlin/Any?): kotlin/Any? // kotlinx.coroutines/CancellableContinuationImpl.tryResume|tryResume(1:0;kotlin.Any?){}[0]
-    open fun tryResume(#A, kotlin/Any?, kotlin/Function1<kotlin/Throwable, kotlin/Unit>?): kotlin/Any? // kotlinx.coroutines/CancellableContinuationImpl.tryResume|tryResume(1:0;kotlin.Any?;kotlin.Function1<kotlin.Throwable,kotlin.Unit>?){}[0]
     open fun tryResumeWithException(kotlin/Throwable): kotlin/Any? // kotlinx.coroutines/CancellableContinuationImpl.tryResumeWithException|tryResumeWithException(kotlin.Throwable){}[0]
     open val callerFrame // kotlinx.coroutines/CancellableContinuationImpl.callerFrame|{}callerFrame[0]
         open fun <get-callerFrame>(): kotlinx.coroutines.internal/CoroutineStackFrame? // kotlinx.coroutines/CancellableContinuationImpl.callerFrame.<get-callerFrame>|<get-callerFrame>(){}[0]
@@ -903,7 +905,7 @@ sealed interface kotlinx.coroutines.selects/SelectClause { // kotlinx.coroutines
     abstract val clauseObject // kotlinx.coroutines.selects/SelectClause.clauseObject|{}clauseObject[0]
         abstract fun <get-clauseObject>(): kotlin/Any // kotlinx.coroutines.selects/SelectClause.clauseObject.<get-clauseObject>|<get-clauseObject>(){}[0]
     abstract val onCancellationConstructor // kotlinx.coroutines.selects/SelectClause.onCancellationConstructor|{}onCancellationConstructor[0]
-        abstract fun <get-onCancellationConstructor>(): kotlin/Function3<kotlinx.coroutines.selects/SelectInstance<*>, kotlin/Any?, kotlin/Any?, kotlin/Function1<kotlin/Throwable, kotlin/Unit>>? // kotlinx.coroutines.selects/SelectClause.onCancellationConstructor.<get-onCancellationConstructor>|<get-onCancellationConstructor>(){}[0]
+        abstract fun <get-onCancellationConstructor>(): kotlin/Function3<kotlinx.coroutines.selects/SelectInstance<*>, kotlin/Any?, kotlin/Any?, kotlin/Function3<kotlin/Throwable, kotlin/Any?, kotlin.coroutines/CoroutineContext, kotlin/Unit>>? // kotlinx.coroutines.selects/SelectClause.onCancellationConstructor.<get-onCancellationConstructor>|<get-onCancellationConstructor>(){}[0]
     abstract val processResFunc // kotlinx.coroutines.selects/SelectClause.processResFunc|{}processResFunc[0]
         abstract fun <get-processResFunc>(): kotlin/Function3<kotlin/Any, kotlin/Any?, kotlin/Any?, kotlin/Any?> // kotlinx.coroutines.selects/SelectClause.processResFunc.<get-processResFunc>|<get-processResFunc>(){}[0]
     abstract val regFunc // kotlinx.coroutines.selects/SelectClause.regFunc|{}regFunc[0]

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

@@ -74,16 +74,21 @@ public interface CancellableContinuation<in T> : Continuation<T> {
     public fun tryResume(value: T, idempotent: Any? = null): Any?
 
     /**
-     * Same as [tryResume] but with [onCancellation] handler that called if and only if the value is not
-     * delivered to the caller because of the dispatch in the process, so that atomicity delivery
-     * guaranteed can be provided by having a cancellation fallback.
+     * Same as [tryResume] but with an [onCancellation] handler that is called if and only if the value is not
+     * delivered to the caller because of the dispatch in the process.
+     *
+     * The purpose of this function is to enable atomic delivery guarantees: either resumption succeeded, passing
+     * the responsibility for [value] to the continuation, or the [onCancellation] block will be invoked,
+     * allowing one to free the resources in [value].
      *
      * Implementation note: current implementation always returns RESUME_TOKEN or `null`
      *
      * @suppress  **This is unstable API and it is subject to change.**
      */
     @InternalCoroutinesApi
-    public fun tryResume(value: T, idempotent: Any?, onCancellation: ((cause: Throwable) -> Unit)?): Any?
+    public fun <R: T> tryResume(
+        value: R, idempotent: Any?, onCancellation: ((cause: Throwable, value: R, context: CoroutineContext) -> Unit)?
+    ): Any?
 
     /**
      * Tries to resume this continuation with the specified [exception] and returns a non-null object token if successful,
@@ -126,7 +131,7 @@ public interface CancellableContinuation<in T> : Continuation<T> {
      * Otherwise, the handler will be invoked as soon as this continuation is cancelled.
      *
      * The installed [handler] should not throw any exceptions.
-     * If it does, they will get caught, wrapped into a [CompletionHandlerException] and
+     * If it does, they will get caught, wrapped into a `CompletionHandlerException` and
      * processed as an uncaught exception in the context of the current coroutine
      * (see [CoroutineExceptionHandler]).
      *
@@ -168,36 +173,60 @@ public interface CancellableContinuation<in T> : Continuation<T> {
     @ExperimentalCoroutinesApi
     public fun CoroutineDispatcher.resumeUndispatchedWithException(exception: Throwable)
 
+    /** @suppress */
+    @Deprecated(
+        "Use the overload that also accepts the `value` and the coroutine context in lambda",
+        level = DeprecationLevel.WARNING,
+        replaceWith = ReplaceWith("resume(value) { cause, _, _ -> onCancellation(cause) }")
+    ) // warning since 1.9.0, was experimental
+    public fun resume(value: T, onCancellation: ((cause: Throwable) -> Unit)?)
+
     /**
-     * Resumes this continuation with the specified `value` and calls the specified `onCancellation`
-     * handler when either resumed too late (when continuation was already cancelled) or, although resumed
-     * successfully (before cancellation), the coroutine's job was cancelled before it had a
-     * chance to run in its dispatcher, so that the suspended function threw an exception
-     * instead of returning this value.
+     * Resumes this continuation with the specified [value], calling the specified [onCancellation] if and only if
+     * the [value] was not successfully used to resume the continuation.
+     *
+     * The [value] can be rejected in two cases (in both of which [onCancellation] will be called):
+     * - Cancellation happened before the handler was resumed;
+     * - The continuation was resumed successfully (before cancellation), but the coroutine's job was cancelled before
+     *   it had a chance to run in its dispatcher, and so the suspended function threw an exception instead of returning
+     *   this value.
      *
      * The installed [onCancellation] handler should not throw any exceptions.
-     * If it does, they will get caught, wrapped into a [CompletionHandlerException] and
+     * If it does, they will get caught, wrapped into a `CompletionHandlerException`, and
      * processed as an 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
-     * code that called the corresponding suspending function, for example:
+     * With this version of [resume], it's possible to pass resources that can not simply be left for the garbage
+     * collector (like file handles, sockets, etc.) and need to be closed explicitly:
      *
      * ```
-     * continuation.resume(resource) {
-     *     resource.close()
+     * continuation.resume(resourceToResumeWith) { _, resourceToClose, _ ->
+     *     resourceToClose.close()
      * }
      * ```
      *
+     * [onCancellation] accepts three arguments:
+     *
+     * - `cause: Throwable` is the exception with which the continuation was cancelled.
+     * - `value` is exactly the same as the [value] passed to [resume] itself.
+     *   In the example above, `resourceToResumeWith` is exactly the same as `resourceToClose`; in particular,
+     *   one could call `resourceToResumeWith.close()` in the lambda for the same effect.
+     *   The reason to reference `resourceToClose` anyway is to avoid a memory allocation due to the lambda
+     *   capturing the `resourceToResumeWith` reference.
+     * - `context` is the [context] of this continuation.
+     *   Like with `value`, the reason this is available as a lambda parameter, even though it is always possible to
+     *   call [context] from the lambda instead, is to allow lambdas to capture less of their environment.
+     *
      * A more complete example and further details are given in
      * the documentation for the [suspendCancellableCoroutine] function.
      *
      * **Note**: The [onCancellation] handler must be fast, non-blocking, and thread-safe.
      * It can be invoked concurrently with the surrounding code.
      * There is no guarantee on the execution context of its invocation.
      */
-    @ExperimentalCoroutinesApi // since 1.2.0
-    public fun resume(value: T, onCancellation: ((cause: Throwable) -> Unit)?)
+    public fun <R: T> resume(
+        value: R, onCancellation: ((cause: Throwable, value: R, context: CoroutineContext) -> Unit)?
+    )
 }
 
 /**
@@ -293,8 +322,10 @@ internal fun <T> CancellableContinuation<T>.invokeOnCancellation(handler: Cancel
  *         override fun onCompleted(resource: T) {
  *             // Resume coroutine with a value provided by the callback and ensure the resource is closed in case
  *             // when the coroutine is cancelled before the caller gets a reference to the resource.
- *             continuation.resume(resource) {
- *                 resource.close() // Close the resource on cancellation
+ *             continuation.resume(resource) { cause, resourceToClose, context ->
+ *                 resourceToClose.close() // Close the resource on cancellation
+ *                 // If we used `resource` instead of `resourceToClose`, this lambda would need to allocate a closure,
+ *                 // but with `resourceToClose`, the lambda does not capture any of its environment.
  *             }
  *         }
  *     // ...