~: [squash] do not postpone explicit cancellation, store Throwable instead of CancelledContinuation in reusability SM, style fixes

Author: qwwdfsad

benchmarks/src/jmh/kotlin/benchmarks/tailcall/SimpleChannel.kt

@@ -83,11 +83,13 @@ class CancellableChannel : SimpleChannel() {
 }
 
 class CancellableReusableChannel : SimpleChannel() {
+    @Suppress("INVISIBLE_MEMBER")
     override suspend fun suspendReceive(): Int = suspendAtomicCancellableCoroutineReusable {
         consumer = it.intercepted()
         COROUTINE_SUSPENDED
     }
 
+    @Suppress("INVISIBLE_MEMBER")
     override suspend fun suspendSend(element: Int) = suspendAtomicCancellableCoroutineReusable<Unit> {
         enqueuedValue = element
         producer = it.intercepted()

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

@@ -239,7 +239,6 @@ internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>):
     if (delegate !is DispatchedContinuation<T>) {
         return CancellableContinuationImpl(delegate, resumeMode = MODE_ATOMIC_DEFAULT)
     }
-
     /*
      * Attempt to claim reusable instance.
      *
@@ -254,9 +253,8 @@ internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>):
      *    thus leaking CC instance for indefinite time.
      * 2) Continuation was cancelled. Then we should prevent any further reuse and bail out.
      */
-    val claimed = delegate.claimReusableCancellableContinuation() ?: return CancellableContinuationImpl(delegate, MODE_ATOMIC_DEFAULT)
-    claimed.resetState()
-    return claimed
+    return delegate.claimReusableCancellableContinuation()?.apply { resetState() }
+        ?: return CancellableContinuationImpl(delegate, MODE_ATOMIC_DEFAULT)
 }
 
 /**

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

@@ -122,10 +122,10 @@ internal open class CancellableContinuationImpl<in T>(
         val completed = isCompleted
         if (resumeMode != MODE_ATOMIC_DEFAULT) return completed // Do not check postponed cancellation for non-reusable continuations
         val dispatched = delegate as? DispatchedContinuation<*> ?: return completed
-        val cancelled = dispatched.checkPostponedCancellation(this) ?: return completed
+        val cause = dispatched.checkPostponedCancellation(this) ?: return completed
         if (!completed) {
             // Note: this cancel may fail if one more concurrent cancel is currently being invoked
-            cancel(cancelled.cause)
+            cancel(cause)
         }
         return true
     }
@@ -148,21 +148,13 @@ internal open class CancellableContinuationImpl<in T>(
     /*
      * Attempt to postpone cancellation for reusable cancellable continuation
      */
-    private fun cancelLater(cancelled: CancelledContinuation): Boolean {
+    private fun cancelLater(cause: Throwable): Boolean {
         if (resumeMode != MODE_ATOMIC_DEFAULT) return false
         val dispatched = (delegate as? DispatchedContinuation<*>) ?: return false
-        return dispatched.postponeCancellation(cancelled)
+        return dispatched.postponeCancellation(cause)
     }
 
     public override fun cancel(cause: Throwable?): Boolean {
-        if (cancelLater(CancelledContinuation(this, cause, handled = state is CancelHandler))) {
-            /*
-             * Here we can't reliably say whether postponed cancellation will be successful, but as it's internal API
-             * and we do not rely on return value, we are free to return `true`.
-             */
-            return true
-        }
-
         _state.loop { state ->
             if (state !is NotCompleted) return false // false if already complete or cancelling
             // Active -- update to final state
@@ -171,12 +163,21 @@ internal open class CancellableContinuationImpl<in T>(
             // Invoke cancel handler if it was present
             if (state is CancelHandler) invokeHandlerSafely { state.invoke(cause) }
             // Complete state update
-            disposeParentHandle()
+            detachChildIfNonResuable()
             dispatchResume(mode = MODE_ATOMIC_DEFAULT)
             return true
         }
     }
 
+    internal fun parentCancelled(cause: Throwable) {
+        /*
+         * Here we can't reliably say whether postponed cancellation will be successful, but as it's internal API
+         * and we do not rely on return value, we are free to return `true`.
+         */
+        if (cancelLater(cause)) return
+        cancel(cause)
+    }
+
     private inline fun invokeHandlerSafely(block: () -> Unit) {
         try {
             block()
@@ -308,7 +309,7 @@ internal open class CancellableContinuationImpl<in T>(
             when (state) {
                 is NotCompleted -> {
                     if (!_state.compareAndSet(state, proposedUpdate)) return@loop // retry on cas failure
-                    disposeParentHandle()
+                    detachChildIfNonResuable()
                     dispatchResume(resumeMode)
                     return null
                 }
@@ -330,14 +331,9 @@ internal open class CancellableContinuationImpl<in T>(
     }
 
     // Unregister from parent job
-    private fun disposeParentHandle() {
-        // If instance is reusable, do not detach on every reuse, #releaseInterceptedContinuation
-        // will do it for us in the end
-        if (isReusable()) return
-        parentHandle?.let { // volatile read parentHandle (once)
-            it.dispose()
-            parentHandle = NonDisposableHandle // release it just in case, to aid GC
-        }
+    private fun detachChildIfNonResuable() {
+        // If instance is reusable, do not detach on every reuse, #releaseInterceptedContinuation will do it for us in the end
+        if (!isReusable()) detachChild()
     }
 
     /**
@@ -361,7 +357,7 @@ internal open class CancellableContinuationImpl<in T>(
                     val update: Any? = if (idempotent == null) value else
                         CompletedIdempotentResult(idempotent, value, state)
                     if (!_state.compareAndSet(state, update)) return@loop // retry on cas failure
-                    disposeParentHandle()
+                    detachChildIfNonResuable()
                     return state
                 }
                 is CompletedIdempotentResult -> {
@@ -383,7 +379,7 @@ internal open class CancellableContinuationImpl<in T>(
                 is NotCompleted -> {
                     val update = CompletedExceptionally(exception)
                     if (!_state.compareAndSet(state, update)) return@loop // retry on cas failure
-                    disposeParentHandle()
+                    detachChildIfNonResuable()
                     return state
                 }
                 else -> return null // cannot resume -- not active anymore

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

@@ -1415,7 +1415,7 @@ internal class ChildContinuation(
     @JvmField val child: CancellableContinuationImpl<*>
 ) : JobCancellingNode<Job>(parent) {
     override fun invoke(cause: Throwable?) {
-        child.cancel(child.getContinuationCancellationCause(job))
+        child.parentCancelled(child.getContinuationCancellationCause(job))
     }
     override fun toString(): String =
         "ChildContinuation[$child]"

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

@@ -47,7 +47,7 @@ internal class DispatchedContinuation<in T>(
      *    }
      *    // state == CC
      *    ```
-     * 4) [CancelledContinuation] continuation was cancelled while being in [suspendAtomicCancellableCoroutineReusable],
+     * 4) [Throwable] continuation was cancelled with this cause while being in [suspendAtomicCancellableCoroutineReusable],
      *    [CancellableContinuationImpl.getResult] will check for cancellation later.
      * 5) [NON_REUSABLE]. CC was cancelled at least once, thus cannot be longer reused.
      *
@@ -84,26 +84,29 @@ internal class DispatchedContinuation<in T>(
          * 3) `NON_REUSABLE` -> nothing, proceed, caller will instantiate CC instance
          */
         _reusableCancellableContinuation.loop { state ->
-            when (state) {
-                null -> {
-                    if (_reusableCancellableContinuation.compareAndSet(null, REUSABLE_CLAIMED)) {
-                        return null
-                    }
+            when {
+                state === null -> {
+                    /*
+                     * null -> CC was not yet published -> we do not compete with cancel
+                     * -> can use plain store instead of CAS
+                     */
+                    _reusableCancellableContinuation.value = REUSABLE_CLAIMED
+                    return null
                 }
-                is CancellableContinuationImpl<*> -> {
+                state is CancellableContinuationImpl<*> -> {
                     if (_reusableCancellableContinuation.compareAndSet(state, REUSABLE_CLAIMED)) {
                         return state as CancellableContinuationImpl<T>
                     }
                 }
-                NON_REUSABLE -> return null
+                state === NON_REUSABLE -> return null
                 else ->  error("Inconsistent state $state")
             }
         }
     }
 
     /**
      * Checks whether there were any attempts to cancel reusable CC while it was in [REUSABLE_CLAIMED] state
-     * and returns [CancelledContinuation] if so, `null` otherwise.
+     * and returns cancellation cause if so, `null` otherwise.
      * If continuation was cancelled, it becomes non-reusable.
      *
      * ```
@@ -115,13 +118,13 @@ internal class DispatchedContinuation<in T>(
      *
      * See [CancellableContinuationImpl.getResult].
      */
-    fun checkPostponedCancellation(continuation: CancellableContinuation<*>): CancelledContinuation? {
+    fun checkPostponedCancellation(continuation: CancellableContinuation<*>): Throwable? {
         _reusableCancellableContinuation.loop { state ->
             // not when(state) to avoid Intrinsics.equals call
             when {
                 state === REUSABLE_CLAIMED -> if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, continuation)) return null
                 state === null -> return null
-                state is CancelledContinuation -> {
+                state is Throwable -> {
                     require(_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE))
                     return state
                 }
@@ -134,7 +137,7 @@ internal class DispatchedContinuation<in T>(
         _reusableCancellableContinuation.loop { state ->
             when (state) {
                 // Do not overwrite cancellation (cancelled CC can't be reused anyway)
-                is CancelledContinuation -> return
+                is Throwable -> return
                 else -> if (_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE)) return
             }
         }
@@ -144,14 +147,14 @@ internal class DispatchedContinuation<in T>(
      * Tries to postpone cancellation if reusable CC is currently in [REUSABLE_CLAIMED] state.
      * Returns `true` if cancellation is (or previously was) postponed, `false` otherwise.
      */
-    fun postponeCancellation(cancelled: CancelledContinuation): Boolean {
+    fun postponeCancellation(cause: Throwable): Boolean {
         _reusableCancellableContinuation.loop { state ->
             when (state) {
                 REUSABLE_CLAIMED -> {
-                    if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, cancelled))
+                    if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, cause))
                         return true
                 }
-                is CancelledContinuation -> return true
+                is Throwable -> return true
                 else -> {
                     if (_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE))
                         return false

kotlinx-coroutines-core/common/src/sync/Mutex.kt

@@ -187,7 +187,7 @@ internal class MutexImpl(locked: Boolean) : Mutex, SelectClause2<Any?, Mutex> {
         return lockSuspend(owner)
     }
 
-    private suspend fun lockSuspend(owner: Any?) = suspendAtomicCancellableCoroutine<Unit> sc@ { cont ->
+    private suspend fun lockSuspend(owner: Any?) = suspendAtomicCancellableCoroutineReusable<Unit> sc@ { cont ->
         val waiter = LockCont(owner, cont)
         _state.loop { state ->
             when (state) {

kotlinx-coroutines-core/common/src/sync/Semaphore.kt

@@ -131,7 +131,7 @@ private class SemaphoreImpl(
         cur + 1
     }
 
-    private suspend fun addToQueueAndSuspend() = suspendAtomicCancellableCoroutine<Unit> sc@ { cont ->
+    private suspend fun addToQueueAndSuspend() = suspendAtomicCancellableCoroutineReusable<Unit> sc@ { cont ->
         val last = this.tail
         val enqIdx = enqIdx.getAndIncrement()
         val segment = getSegment(last, enqIdx / SEGMENT_SIZE)

kotlinx-coroutines-core/jvm/test/ReusableCancellableContinuationTest.kt

@@ -60,14 +60,30 @@ class ReusableCancellableContinuationTest : TestBase() {
     }
 
     @Test
-    fun testNotCancelledOnClaimedResume() = runTest {
+    fun testNotCancelledOnClaimedResume() = runTest({ it is CancellationException }) {
         expect(1)
+        // Bind child at first
+        var continuation: Continuation<*>? = null
         suspendAtomicCancellableCoroutineReusable<Unit> {
-            it.cancel()
-            it.resume(Unit)
+            expect(2)
+            continuation = it
+            launch {
+                expect(3)
+                it.resume(Unit)
+            }
         }
+        println("1")
+        expect(4)
         ensureActive()
-        finish(2)
+        // Verify child was bound
+        assertNotNull(FieldWalker.walk(coroutineContext[Job]!!).single { it === continuation })
+        suspendAtomicCancellableCoroutineReusable<Unit> {
+            expect(5)
+            coroutineContext[Job]!!.cancel()
+            it.resume(Unit)
+        }
+        assertFalse(isActive)
+        finish(6)
     }
 
     @Test