Rework CompletionHandler to avoid subclassing a functional type (#4010)

Workaround for KT-64075

Also, thoroughly document the contracts imposed on the callbacks provided to `invokeOnCompletion` and `invokeOnCancellation`.

Author: dkhalanskyjb

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

@@ -87,7 +87,6 @@ public final class kotlinx/coroutines/CancellableContinuationKt {
 public final class kotlinx/coroutines/ChildContinuation {
 	public final field child Lkotlinx/coroutines/CancellableContinuationImpl;
 	public fun <init> (Lkotlinx/coroutines/CancellableContinuationImpl;)V
-	public synthetic fun invoke (Ljava/lang/Object;)Ljava/lang/Object;
 	public fun invoke (Ljava/lang/Throwable;)V
 }
 
@@ -1332,12 +1331,11 @@ public abstract interface class kotlinx/coroutines/selects/SelectClause1 : kotli
 public abstract interface class kotlinx/coroutines/selects/SelectClause2 : kotlinx/coroutines/selects/SelectClause {
 }
 
-public class kotlinx/coroutines/selects/SelectImplementation : kotlinx/coroutines/selects/SelectBuilder, kotlinx/coroutines/selects/SelectInstanceInternal {
+public class kotlinx/coroutines/selects/SelectImplementation : kotlinx/coroutines/CancelHandler, kotlinx/coroutines/selects/SelectBuilder, kotlinx/coroutines/selects/SelectInstanceInternal {
 	public fun <init> (Lkotlin/coroutines/CoroutineContext;)V
 	public fun disposeOnCompletion (Lkotlinx/coroutines/DisposableHandle;)V
 	public fun doSelect (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun getContext ()Lkotlin/coroutines/CoroutineContext;
-	public synthetic fun invoke (Ljava/lang/Object;)Ljava/lang/Object;
 	public fun invoke (Ljava/lang/Throwable;)V
 	public fun invoke (Lkotlinx/coroutines/selects/SelectClause0;Lkotlin/jvm/functions/Function1;)V
 	public fun invoke (Lkotlinx/coroutines/selects/SelectClause1;Lkotlin/jvm/functions/Function2;)V

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

@@ -67,7 +67,7 @@ private class AwaitAll<T>(private val deferreds: Array<out Deferred<T>>) {
             val deferred = deferreds[i]
             deferred.start() // To properly await lazily started deferreds
             AwaitAllNode(cont).apply {
-                handle = deferred.invokeOnCompletion(asHandler)
+                handle = deferred.invokeOnCompletion(handler = this)
             }
         }
         val disposer = DisposeHandlersOnCancel(nodes)
@@ -79,11 +79,11 @@ private class AwaitAll<T>(private val deferreds: Array<out Deferred<T>>) {
             // it is already complete while handlers were being installed -- dispose them all
             disposer.disposeAll()
         } else {
-            cont.invokeOnCancellation(handler = disposer.asHandler)
+            cont.invokeOnCancellation(handler = disposer)
         }
     }
 
-    private inner class DisposeHandlersOnCancel(private val nodes: Array<AwaitAllNode>) : CancelHandler() {
+    private inner class DisposeHandlersOnCancel(private val nodes: Array<AwaitAllNode>) : CancelHandler {
         fun disposeAll() {
             nodes.forEach { it.handle.dispose() }
         }

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

@@ -122,28 +122,27 @@ public interface CancellableContinuation<in T> : Continuation<T> {
 
     /**
      * Registers a [handler] to be **synchronously** invoked on [cancellation][cancel] (regular or exceptional) of this continuation.
-     * When the continuation is already cancelled, the handler is immediately invoked
-     * with the cancellation exception. Otherwise, the handler will be invoked as soon as this
-     * continuation is cancelled.
+     * When the continuation is already cancelled, the handler is immediately invoked with the cancellation exception.
+     * 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
      * processed as an uncaught exception in the context of the current coroutine
      * (see [CoroutineExceptionHandler]).
      *
-     * At most one [handler] can be installed on a continuation. Attempt to call `invokeOnCancellation` second
-     * time produces [IllegalStateException].
+     * At most one [handler] can be installed on a continuation.
+     * Attempting to call `invokeOnCancellation` a second time produces an [IllegalStateException].
      *
      * This handler is also called when this continuation [resumes][Continuation.resume] normally (with a value) and then
      * is cancelled while waiting to be dispatched. More generally speaking, this handler is called whenever
      * the caller of [suspendCancellableCoroutine] is getting a [CancellationException].
      *
-     * A typical example for `invokeOnCancellation` usage is given in
+     * A typical example of `invokeOnCancellation` usage is given in
      * the documentation for the [suspendCancellableCoroutine] function.
      *
-     * **Note**: Implementation of `CompletionHandler` must be fast, non-blocking, and thread-safe.
-     * This `handler` can be invoked concurrently with the surrounding code.
-     * There is no guarantee on the execution context in which the `handler` will be invoked.
+     * **Note**: Implementations of [CompletionHandler] must be fast, non-blocking, and thread-safe.
+     * This [handler] can be invoked concurrently with the surrounding code.
+     * There is no guarantee on the execution context in which the [handler] will be invoked.
      */
     public fun invokeOnCancellation(handler: CompletionHandler)
 
@@ -201,6 +200,15 @@ public interface CancellableContinuation<in T> : Continuation<T> {
     public fun resume(value: T, onCancellation: ((cause: Throwable) -> Unit)?)
 }
 
+/**
+ * A version of `invokeOnCancellation` that accepts a class as a handler instead of a lambda, but identical otherwise.
+ * This allows providing a custom [toString] instance that will look better during debugging.
+ */
+internal fun <T> CancellableContinuation<T>.invokeOnCancellation(handler: CancelHandler) = when (this) {
+    is CancellableContinuationImpl -> invokeOnCancellationInternal(handler)
+    else -> throw UnsupportedOperationException("third-party implementation of CancellableContinuation is not supported")
+}
+
 /**
  * Suspends the coroutine like [suspendCoroutine], but providing a [CancellableContinuation] to
  * the [block]. This function throws a [CancellationException] if the [Job] of the coroutine is
@@ -373,9 +381,9 @@ internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>):
  */
 @InternalCoroutinesApi
 public fun CancellableContinuation<*>.disposeOnCancellation(handle: DisposableHandle): Unit =
-    invokeOnCancellation(handler = DisposeOnCancel(handle).asHandler)
+    invokeOnCancellation(handler = DisposeOnCancel(handle))
 
-private class DisposeOnCancel(private val handle: DisposableHandle) : CancelHandler() {
+private class DisposeOnCancel(private val handle: DisposableHandle) : CancelHandler {
     override fun invoke(cause: Throwable?) = handle.dispose()
     override fun toString(): String = "DisposeOnCancel[$handle]"
 }

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

@@ -234,12 +234,12 @@ internal open class CancellableContinuationImpl<in T>(
         }
     }
 
-    private fun callCancelHandler(handler: CompletionHandler, cause: Throwable?) =
+    private fun callCancelHandler(handler: InternalCompletionHandler, cause: Throwable?) =
         /*
         * :KLUDGE: We have to invoke a handler in platform-specific way via `invokeIt` extension,
         * because we play type tricks on Kotlin/JS and handler is not necessarily a function there
         */
-        callCancelHandlerSafely { handler.invokeIt(cause) }
+        callCancelHandlerSafely { handler.invoke(cause) }
 
     fun callCancelHandler(handler: CancelHandler, cause: Throwable?) =
         callCancelHandlerSafely { handler.invoke(cause) }
@@ -343,7 +343,7 @@ internal open class CancellableContinuationImpl<in T>(
         // Install the handle
         val handle = parent.invokeOnCompletion(
             onCancelling = true,
-            handler = ChildContinuation(this).asHandler
+            handler = ChildContinuation(this)
         )
         _parentHandle.compareAndSet(null, handle)
         return handle
@@ -390,10 +390,9 @@ internal open class CancellableContinuationImpl<in T>(
         invokeOnCancellationImpl(segment)
     }
 
-    public override fun invokeOnCancellation(handler: CompletionHandler) {
-        val cancelHandler = makeCancelHandler(handler)
-        invokeOnCancellationImpl(cancelHandler)
-    }
+    override fun invokeOnCancellation(handler: CompletionHandler) = invokeOnCancellation(CancelHandler.UserSupplied(handler))
+
+    internal fun invokeOnCancellationInternal(handler: CancelHandler) = invokeOnCancellationImpl(handler)
 
     private fun invokeOnCancellationImpl(handler: Any) {
         assert { handler is CancelHandler || handler is Segment<*> }
@@ -461,9 +460,6 @@ internal open class CancellableContinuationImpl<in T>(
         error("It's prohibited to register multiple handlers, tried to register $handler, already has $state")
     }
 
-    private fun makeCancelHandler(handler: CompletionHandler): CancelHandler =
-        if (handler is CancelHandler) handler else InvokeOnCancel(handler)
-
     private fun dispatchResume(mode: Int) {
         if (tryResume()) return // completed before getResult invocation -- bail out
         // otherwise, getResult has already commenced, i.e. completed later or in other thread
@@ -625,19 +621,46 @@ private object Active : NotCompleted {
 }
 
 /**
- * Base class for all [CancellableContinuation.invokeOnCancellation] handlers to avoid an extra instance
- * on JVM, yet support JS where you cannot extend from a functional type.
+ * Essentially the same as just a function from `Throwable?` to `Unit`.
+ * The only thing implementors can do is call [invoke].
+ * The reason this abstraction exists is to allow providing a readable [toString] in the list of completion handlers
+ * as seen from the debugger.
+ * Use [UserSupplied] to create an instance from a lambda.
+ * We can't avoid defining a separate type, because on JS, you can't inherit from a function type.
+ *
+ * @see InternalCompletionHandler for a very similar interface, but used for handling completion and not cancellation.
  */
-internal abstract class CancelHandler : CancelHandlerBase(), NotCompleted
-
-// Wrapper for lambdas, for the performance sake CancelHandler can be subclassed directly
-private class InvokeOnCancel( // Clashes with InvokeOnCancellation
-    private val handler: CompletionHandler
-) : CancelHandler() {
-    override fun invoke(cause: Throwable?) {
-        handler.invoke(cause)
+internal interface CancelHandler : NotCompleted {
+    /**
+     * Signals cancellation.
+     *
+     * This function:
+     * - Does not throw any exceptions.
+     *   Violating this rule in an implementation leads to [handleUncaughtCoroutineException] being called with a
+     *   [CompletionHandlerException] wrapping the thrown exception.
+     * - Is fast, non-blocking, and thread-safe.
+     * - Can be invoked concurrently with the surrounding code.
+     * - Can be invoked from any context.
+     *
+     * The meaning of `cause` that is passed to the handler is:
+     * - It is `null` if the continuation was cancelled directly via [CancellableContinuation.cancel] without a `cause`.
+     * - It is an instance of [CancellationException] if the continuation was _normally_ cancelled from the outside.
+     *   **It should not be treated as an error**. In particular, it should not be reported to error logs.
+     * - Otherwise, the continuation had cancelled with an _error_.
+     */
+    fun invoke(cause: Throwable?)
+
+    /**
+     * A lambda passed from outside the coroutine machinery.
+     *
+     * See the requirements for [CancelHandler.invoke] when implementing this function.
+     */
+    class UserSupplied(private val handler: (cause: Throwable?) -> Unit) : CancelHandler {
+        /** @suppress */
+        override fun invoke(cause: Throwable?) { handler(cause) }
+
+        override fun toString() = "CancelHandler.UserSupplied[${handler.classSimpleName}@$hexAddress]"
     }
-    override fun toString() = "InvokeOnCancel[${handler.classSimpleName}@$hexAddress]"
 }
 
 // Completed with additional metadata

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

@@ -1,43 +1,71 @@
 package kotlinx.coroutines
 
-import kotlinx.coroutines.internal.*
-
 /**
  * Handler for [Job.invokeOnCompletion] and [CancellableContinuation.invokeOnCancellation].
  *
- * 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.
- *
- * The meaning of `cause` that is passed to the handler:
- * - Cause is `null` when the job has completed normally.
- * - Cause is an instance of [CancellationException] when the job was cancelled _normally_.
+ * The meaning of `cause` that is passed to the handler is:
+ * - It is `null` if the job has completed normally or the continuation was cancelled without a `cause`.
+ * - It is an instance of [CancellationException] if the job or the continuation was cancelled _normally_.
  *   **It should not be treated as an error**. In particular, it should not be reported to error logs.
- * - Otherwise, the job had _failed_.
+ * - Otherwise, the job or the continuation had _failed_.
+ *
+ * A function used for this should not throw any exceptions.
+ * If it does, they will get caught, wrapped into [CompletionHandlerException], and then either
+ * - passed to [handleCoroutineException] for [CancellableContinuation.invokeOnCancellation]
+ *   and, for [Job] instances that are coroutines, [Job.invokeOnCompletion], or
+ * - for [Job] instances that are not coroutines, simply thrown, potentially crashing unrelated code.
+ *
+ * Functions used for this must be fast, non-blocking, and thread-safe.
+ * This handler can be invoked concurrently with the surrounding code.
+ * There is no guarantee on the execution context in which the function is invoked.
  *
  * **Note**: This type 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 type should not be used in general application code.
- * Implementations of `CompletionHandler` must be fast and _lock-free_.
  */
+// TODO: deprecate. This doesn't seem better than a simple function type.
 public typealias CompletionHandler = (cause: Throwable?) -> Unit
 
-// We want class that extends LockFreeLinkedListNode & CompletionHandler but we cannot do it on Kotlin/JS,
-// so this expect class provides us with the corresponding abstraction in a platform-agnostic way.
-internal expect abstract class CompletionHandlerBase() : LockFreeLinkedListNode {
-    abstract fun invoke(cause: Throwable?)
-}
+/**
+ * Essentially the same as just a function from `Throwable?` to `Unit`.
+ * The only thing implementors can do is call [invoke].
+ * The reason this abstraction exists is to allow providing a readable [toString] in the list of completion handlers
+ * as seen from the debugger.
+ * Use [UserSupplied] to create an instance from a lambda.
+ * We can't avoid defining a separate type, because on JS, you can't inherit from a function type.
+ *
+ * @see CancelHandler for a very similar interface, but used for handling cancellation and not completion.
+ */
+internal interface InternalCompletionHandler {
+    /**
+     * Signals completion.
+     *
+     * This function:
+     * - Does not throw any exceptions.
+     *   For [Job] instances that are coroutines, exceptions thrown by this function will be caught, wrapped into
+     *   [CompletionHandlerException], and passed to [handleCoroutineException], but for those that are not coroutines,
+     *   they will just be rethrown, potentially crashing unrelated code.
+     * - Is fast, non-blocking, and thread-safe.
+     * - Can be invoked concurrently with the surrounding code.
+     * - Can be invoked from any context.
+     *
+     * The meaning of `cause` that is passed to the handler is:
+     * - It is `null` if the job has completed normally.
+     * - It is an instance of [CancellationException] if the job was cancelled _normally_.
+     *   **It should not be treated as an error**. In particular, it should not be reported to error logs.
+     * - Otherwise, the job had _failed_.
+     */
+    fun invoke(cause: Throwable?)
 
-internal expect val CompletionHandlerBase.asHandler: CompletionHandler
+    /**
+     * A lambda passed from outside the coroutine machinery.
+     *
+     * See the requirements for [InternalCompletionHandler.invoke] when implementing this function.
+     */
+    class UserSupplied(private val handler: (cause: Throwable?) -> Unit) : InternalCompletionHandler {
+        /** @suppress */
+        override fun invoke(cause: Throwable?) { handler(cause) }
 
-// More compact version of CompletionHandlerBase for CancellableContinuation with same workaround for JS
-internal expect abstract class CancelHandlerBase() {
-    abstract fun invoke(cause: Throwable?)
+        override fun toString() = "InternalCompletionHandler.UserSupplied[${handler.classSimpleName}@$hexAddress]"
+    }
 }
-
-internal expect val CancelHandlerBase.asHandler: CompletionHandler
-
-// :KLUDGE: We have to invoke a handler in platform-specific way via `invokeIt` extension,
-// because we play type tricks on Kotlin/JS and handler is not necessarily a function there
-internal expect fun CompletionHandler.invokeIt(cause: Throwable?)
-
-internal inline fun <reified T> CompletionHandler.isHandlerOf(): Boolean = this is T

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

@@ -1,7 +1,7 @@
 package kotlinx.coroutines
 
 /**
- * This exception gets thrown if an exception is caught while processing [CompletionHandler] invocation for [Job].
+ * This exception gets thrown if an exception is caught while processing [InternalCompletionHandler] invocation for [Job].
  *
  * @suppress **This an internal API and should not be used from general code.**
  */