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

Author: qwwdfsad

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

@@ -253,7 +253,7 @@ 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.
      */
-    return delegate.claimReusableCancellableContinuation()?.apply { resetState() }
+    return delegate.claimReusableCancellableContinuation()?.takeIf { it.resetState() }
         ?: return CancellableContinuationImpl(delegate, MODE_ATOMIC_DEFAULT)
 }
 

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

@@ -78,22 +78,23 @@ internal open class CancellableContinuationImpl<in T>(
         // This method does nothing. Leftover for binary compatibility with old compiled code
     }
 
-    private fun invalidateReusability() {
-        if (resumeMode != MODE_ATOMIC_DEFAULT) return
-        (delegate as? DispatchedContinuation<*>)?.makeCancellationNonReusable()
-    }
-
     private fun isReusable(): Boolean = delegate is DispatchedContinuation<*> && delegate.isReusable
 
     /**
      * Resets cancellability state in order to [suspendAtomicCancellableCoroutineReusable] to work.
      * Invariant: used only by [suspendAtomicCancellableCoroutineReusable] in [REUSABLE_CLAIMED] state.
      */
-    internal fun resetState() {
+    internal fun resetState(): Boolean {
         assert { parentHandle !== NonDisposableHandle }
-        assert { _state.value !is NotCompleted }
+        val state = _state.value
+        assert { state !is NotCompleted }
+        if (state is CompletedIdempotentResult) {
+            detachChild()
+            return false
+        }
         _decision.value = UNDECIDED
         _state.value = Active
+        return true
     }
 
     /**
@@ -170,12 +171,10 @@ internal open class CancellableContinuationImpl<in T>(
     }
 
     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)
+        // Even if cancellation has failed, we should detach child to avoid potential leak
+        detachChildIfNonResuable()
     }
 
     private inline fun invokeHandlerSafely(block: () -> Unit) {
@@ -347,10 +346,6 @@ internal open class CancellableContinuationImpl<in T>(
     }
 
     override fun tryResume(value: T, idempotent: Any?): Any? {
-        // Important invariant: if idempotent tryResume is used, then regular resume cannot be used, since the
-        // thread that invokes this tryResume install its descriptor into the pointer to the corresponding node first,
-        // so no other thread can find it and invoke resume.
-        if (idempotent != null) invalidateReusability()
         _state.loop { state ->
             when (state) {
                 is NotCompleted -> {

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

@@ -11,12 +11,6 @@ import kotlin.jvm.*
 
 @SharedImmutable
 private val UNDEFINED = Symbol("UNDEFINED")
-
-// Internal to avoid synthetic accessors
-@SharedImmutable
-@JvmField
-internal val NON_REUSABLE = Symbol("NON_REUSABLE")
-
 @SharedImmutable
 @JvmField
 internal val REUSABLE_CLAIMED = Symbol("REUSABLE_CLAIMED")
@@ -36,7 +30,8 @@ internal class DispatchedContinuation<in T>(
     /**
      * Possible states of reusability:
      *
-     * 1) `null`. Cancellable continuation wasn't yet attempted to be reused.
+     * 1) `null`. Cancellable continuation wasn't yet attempted to be reused or
+     *     way used and then invalidated (e.g. because of the cancellation).
      * 2) [CancellableContinuation]. Continuation to be/that is being reused.
      * 3) [REUSABLE_CLAIMED]. CC is currently being reused and its owner executes `suspend` block:
      *    ```
@@ -49,7 +44,6 @@ internal class DispatchedContinuation<in T>(
      *    ```
      * 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.
      *
      * [REUSABLE_CLAIMED] state is required to prevent the lost resume in the channel.
      * AbstractChannel.receive method relies on the fact that the following pattern
@@ -69,7 +63,7 @@ internal class DispatchedContinuation<in T>(
         get() = _reusableCancellableContinuation.value as? CancellableContinuationImpl<*>
 
     public val isReusable: Boolean
-        get() = _reusableCancellableContinuation.value.let { it != null && it != NON_REUSABLE }
+        get() = _reusableCancellableContinuation.value != null
 
     /**
      * Claims the continuation for [suspendAtomicCancellableCoroutineReusable] block,
@@ -81,7 +75,6 @@ internal class DispatchedContinuation<in T>(
          * Transitions:
          * 1) `null` -> claimed, caller will instantiate CC instance
          * 2) `CC` -> claimed, caller will reuse CC instance
-         * 3) `NON_REUSABLE` -> nothing, proceed, caller will instantiate CC instance
          */
         _reusableCancellableContinuation.loop { state ->
             when {
@@ -98,7 +91,6 @@ internal class DispatchedContinuation<in T>(
                         return state as CancellableContinuationImpl<T>
                     }
                 }
-                state === NON_REUSABLE -> return null
                 else ->  error("Inconsistent state $state")
             }
         }
@@ -125,24 +117,14 @@ internal class DispatchedContinuation<in T>(
                 state === REUSABLE_CLAIMED -> if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, continuation)) return null
                 state === null -> return null
                 state is Throwable -> {
-                    require(_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE))
+                    require(_reusableCancellableContinuation.compareAndSet(state, null))
                     return state
                 }
                 else -> return null // Is not reusable
             }
         }
     }
 
-    fun makeCancellationNonReusable() {
-        _reusableCancellableContinuation.loop { state ->
-            when (state) {
-                // Do not overwrite cancellation (cancelled CC can't be reused anyway)
-                is Throwable -> return
-                else -> if (_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE)) return
-            }
-        }
-    }
-
     /**
      * Tries to postpone cancellation if reusable CC is currently in [REUSABLE_CLAIMED] state.
      * Returns `true` if cancellation is (or previously was) postponed, `false` otherwise.
@@ -156,7 +138,8 @@ internal class DispatchedContinuation<in T>(
                 }
                 is Throwable -> return true
                 else -> {
-                    if (_reusableCancellableContinuation.compareAndSet(state, NON_REUSABLE))
+                    // Invalidate
+                    if (_reusableCancellableContinuation.compareAndSet(state, null))
                         return false
                 }
             }

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

@@ -5,6 +5,7 @@
 package kotlinx.coroutines
 
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.selects.*
 import org.junit.Test
 import kotlin.coroutines.*
 import kotlin.test.*
@@ -67,12 +68,11 @@ class ReusableCancellableContinuationTest : TestBase() {
         suspendAtomicCancellableCoroutineReusable<Unit> {
             expect(2)
             continuation = it
-            launch {
+            launch {  // Attach to the parent, avoid fast path
                 expect(3)
                 it.resume(Unit)
             }
         }
-        println("1")
         expect(4)
         ensureActive()
         // Verify child was bound
@@ -105,7 +105,7 @@ class ReusableCancellableContinuationTest : TestBase() {
     fun testResumeRegularDoesntPreservesReference() = runTest {
         expect(1)
         var cont: Continuation<Unit>? = null
-        launch {
+        launch { // Attach to the parent, avoid fast path
             cont!!.resumeWith(Result.success(Unit))
         }
         suspendAtomicCancellableCoroutine<Unit> {
@@ -131,4 +131,65 @@ class ReusableCancellableContinuationTest : TestBase() {
             finish(2)
         }
     }
+
+    @Test
+    fun testPropagatedCancel() = runTest({it is CancellationException}) {
+        val currentJob = coroutineContext[Job]!!
+        expect(1)
+        // Bind child at first
+        suspendAtomicCancellableCoroutineReusable<Unit> {
+            expect(2)
+            // Attach to the parent, avoid fast path
+            launch {
+                expect(3)
+                it.resume(Unit)
+            }
+        }
+        expect(4)
+        ensureActive()
+        // Verify child was bound
+        assertEquals(1, FieldWalker.walk(currentJob).count { it is CancellableContinuation<*> })
+        currentJob.cancel()
+        assertFalse(isActive)
+        // Child detached
+        assertEquals(0, FieldWalker.walk(currentJob).count { it is CancellableContinuation<*> })
+        suspendAtomicCancellableCoroutineReusable<Unit> { it.resume(Unit) }
+        suspendAtomicCancellableCoroutineReusable<Unit> { it.resume(Unit) }
+        assertEquals(0, FieldWalker.walk(currentJob).count { it is CancellableContinuation<*> })
+
+        try {
+            suspendAtomicCancellableCoroutineReusable<Unit> {}
+        } catch (e: CancellationException) {
+            assertEquals(0, FieldWalker.walk(currentJob).count { it is CancellableContinuation<*> })
+            finish(5)
+        }
+    }
+
+    @Test
+    fun testChannelMemoryLeak() = runTest {
+        val iterations = 100
+        val channel = Channel<Unit>()
+        launch {
+            repeat(iterations) {
+                select {
+                    channel.onSend(Unit) {}
+                }
+            }
+        }
+
+        val receiver = launch {
+            repeat(iterations) {
+                channel.receive()
+            }
+            expect(2)
+            val job = coroutineContext[Job]!!
+            // 1 for reusable CC, another one for outer joiner
+            assertEquals(2, FieldWalker.walk(job).count { it is CancellableContinuation<*> })
+        }
+        expect(1)
+        receiver.join()
+        // Reference should be claimed at this point
+        assertEquals(0, FieldWalker.walk(receiver).count { it is CancellableContinuation<*> })
+        finish(3)
+    }
 }