~ Rollback manual continuation dispose on suspend block exception

Author: elizarov

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

@@ -80,8 +80,10 @@ public class kotlinx/coroutines/CancellableContinuationImpl : kotlin/coroutines/
 
 public final class kotlinx/coroutines/CancellableContinuationKt {
 	public static final fun disposeOnCancellation (Lkotlinx/coroutines/CancellableContinuation;Lkotlinx/coroutines/DisposableHandle;)V
+	public static final fun suspendAtomicCancellableCoroutine (Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun suspendAtomicCancellableCoroutine (ZLkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun suspendAtomicCancellableCoroutine$default (ZLkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
+	public static final fun suspendCancellableCoroutine (Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
 public abstract interface class kotlinx/coroutines/ChildHandle : kotlinx/coroutines/DisposableHandle {
@@ -501,11 +503,6 @@ public final class kotlinx/coroutines/SupervisorKt {
 	public static final fun supervisorScope (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
-public final class kotlinx/coroutines/SuspendCancellableCoroutineKt {
-	public static final fun suspendAtomicCancellableCoroutine (Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-	public static final fun suspendCancellableCoroutine (Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-}
-
 public abstract interface class kotlinx/coroutines/ThreadContextElement : kotlin/coroutines/CoroutineContext$Element {
 	public abstract fun restoreThreadContext (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Object;)V
 	public abstract fun updateThreadContext (Lkotlin/coroutines/CoroutineContext;)Ljava/lang/Object;

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

@@ -295,7 +295,7 @@ internal class DispatchedCoroutine<in T>(
     fun getResult(): Any? {
         if (trySuspend()) return COROUTINE_SUSPENDED
         // When scope coroutine does not suspend on Kotlin/Native it shall dispose its continuation which it will not use
-        disposeContinuation { uCont }
+//        disposeContinuation { uCont }
         // otherwise, onCompletionInternal was already invoked & invoked tryResume, and the result is in the state
         val state = this.state.unboxState()
         if (state is CompletedExceptionally) throw state.cause

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

@@ -6,6 +6,7 @@ package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
 
 // --------------- cancellable continuations ---------------
 
@@ -182,9 +183,20 @@ public interface CancellableContinuation<in T> : Continuation<T> {
  * Suspends the coroutine like [suspendCoroutine], but providing a [CancellableContinuation] to
  * the [block]. This function throws a [CancellationException] if the coroutine is cancelled or completed while suspended.
  */
-public expect suspend inline fun <T> suspendCancellableCoroutine(
+public suspend inline fun <T> suspendCancellableCoroutine(
     crossinline block: (CancellableContinuation<T>) -> Unit
-): T
+): T =
+    suspendCoroutineUninterceptedOrReturn { uCont ->
+        val cancellable = CancellableContinuationImpl(uCont.intercepted(), resumeMode = MODE_CANCELLABLE)
+        /*
+         * For non-atomic cancellation we setup parent-child relationship immediately
+         * in case when `block` blocks the current thread (e.g. Rx2 with trampoline scheduler), but
+         * properly supports cancellation.
+         */
+        cancellable.initCancellability()
+        block(cancellable)
+        cancellable.getResult()
+    }
 
 /**
  * Suspends the coroutine like [suspendCancellableCoroutine], but with *atomic cancellation*.
@@ -197,19 +209,31 @@ public expect suspend inline fun <T> suspendCancellableCoroutine(
  * @suppress **This an internal API and should not be used from general code.**
  */
 @InternalCoroutinesApi
-public expect suspend inline fun <T> suspendAtomicCancellableCoroutine(
+public suspend inline fun <T> suspendAtomicCancellableCoroutine(
     crossinline block: (CancellableContinuation<T>) -> Unit
-): T
+): T =
+    suspendCoroutineUninterceptedOrReturn { uCont ->
+        val cancellable = CancellableContinuationImpl(uCont.intercepted(), resumeMode = MODE_ATOMIC_DEFAULT)
+        block(cancellable)
+        cancellable.getResult()
+    }
 
 /**
  *  Suspends coroutine similar to [suspendAtomicCancellableCoroutine], but an instance of [CancellableContinuationImpl] is reused if possible.
  */
-internal expect suspend inline fun <T> suspendAtomicCancellableCoroutineReusable(
+internal suspend inline fun <T> suspendAtomicCancellableCoroutineReusable(
     crossinline block: (CancellableContinuation<T>) -> Unit
-): T
+): T = suspendCoroutineUninterceptedOrReturn { uCont ->
+    val cancellable = getOrCreateCancellableContinuation(uCont.intercepted())
+    block(cancellable)
+    cancellable.getResult()
+}
 
 internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>): CancellableContinuationImpl<T> {
-    if (delegate !is DispatchedContinuation<T>) {
+    // If used outside of our dispatcher
+    // NOTE: Reuse is not support on Kotlin/Native due to platform peculiarities making it had to properly
+    //       split DispatchedContinuation / CancellableContinuationImpl state across workers.
+    if (!isReuseSupportedInPlatform() || delegate !is DispatchedContinuation<T> ) {
         return CancellableContinuationImpl(delegate, resumeMode = MODE_ATOMIC_DEFAULT)
     }
     /*

kotlinx-coroutines-core/common/src/internal/Sharing.common.kt

@@ -18,7 +18,6 @@ internal expect fun <T> Continuation<T>.asLocalOrNullIfNotUsed() : Continuation<
 internal expect fun <T> Continuation<T>.useLocal() : Continuation<T>
 internal expect fun <T> Continuation<T>.shareableInterceptedResumeCancellableWith(result: Result<T>)
 internal expect fun <T> Continuation<T>.shareableInterceptedResumeWith(result: Result<T>)
-internal expect fun <T> Continuation<T>.shareableDispose()
 internal expect fun disposeContinuation(cont: () -> Continuation<*>)
 internal expect fun <T> CancellableContinuationImpl<T>.shareableResume(delegate: Continuation<T>, useMode: Int)
 

kotlinx-coroutines-core/js/src/SuspendCancellableCoroutine.kt

@@ -45,9 +45,6 @@ internal actual inline fun <T> Continuation<T>.shareableInterceptedResumeWith(re
     intercepted().resumeWith(result)
 }
 
-@Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun <T> Continuation<T>.shareableDispose() {}
-
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
 internal actual inline fun disposeContinuation(cont: () -> Continuation<*>) {}
 

kotlinx-coroutines-core/js/src/internal/Sharing.kt

@@ -58,10 +58,6 @@ internal actual inline fun <T> Continuation<T>.shareableInterceptedResumeWith(re
     intercepted().resumeWith(result)
 }
 
-@InlineOnly
-@Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun <T> Continuation<T>.shareableDispose() {}
-
 @InlineOnly
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
 internal actual inline fun disposeContinuation(cont: () -> Continuation<*>) {}

kotlinx-coroutines-core/jvm/src/SuspendCancellableCoroutine.kt

@@ -78,11 +78,6 @@ internal actual fun <T> Continuation<T>.shareableInterceptedResumeWith(result: R
     }
 }
 
-internal actual inline fun <T> Continuation<T>.shareableDispose() {
-    this as ShareableContinuation<T> // must have been shared
-    disposeRef()
-}
-
 internal actual fun <T, R> (suspend (T) -> R).asShareable(): suspend (T) -> R =
     ShareableBlock(this)