Detect suspendCancellableCoroutine right after suspendCancellableCoro…

…utineReusable within the same state machine and properly cleanup its child handle when its block completes Fixes #1855
Kotlin · Mar 11, 2020 · fc46a98 · fc46a98
1 parent 262fb66
commit fc46a98
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/kotlinx-coroutines-core/common/src/CancellableContinuationImpl.kt b/kotlinx-coroutines-core/common/src/CancellableContinuationImpl.kt
@@ -85,12 +85,13 @@ internal open class CancellableContinuationImpl<in T>(
         // This method does nothing. Leftover for binary compatibility with old compiled code
     }
 
-    private fun isReusable(): Boolean = delegate is DispatchedContinuation<*> && delegate.isReusable
+    private fun isReusable(): Boolean = delegate is DispatchedContinuation<*> && delegate.isReusable(this)
 
     /**
      * Resets cancellability state in order to [suspendAtomicCancellableCoroutineReusable] to work.
      * Invariant: used only by [suspendAtomicCancellableCoroutineReusable] in [REUSABLE_CLAIMED] state.
      */
+    @JvmName("resetState") // Prettier stack traces
     internal fun resetState(): Boolean {
         assert { parentHandle !== NonDisposableHandle }
         val state = _state.value

diff --git a/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt b/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt
@@ -63,8 +63,24 @@ internal class DispatchedContinuation<in T>(
     public val reusableCancellableContinuation: CancellableContinuationImpl<*>?
         get() = _reusableCancellableContinuation.value as? CancellableContinuationImpl<*>
 
-    public val isReusable: Boolean
-        get() = _reusableCancellableContinuation.value != null
+    public fun isReusable(requester: CancellableContinuationImpl<*>): Boolean {
+        /*
+         * Reusability control:
+         * `null` -> no reusability at all, false
+         * If current state is not CCI, then we are within `suspendAtomicCancellableCoroutineReusable`, true
+         * Else, if result is CCI === requester.
+         * Identity check my fail for the following pattern:
+         * ```
+         * loop:
+         * suspendAtomicCancellableCoroutineReusable { } // Reusable, outer coroutine stores the child handle
+         * suspendCancellableCoroutine { } // **Not reusable**, handle should be disposed after {}, otherwise
+         * it will leak because it won't be freed by `releaseInterceptedContinuation`
+         * ```
+         */
+        val value = _reusableCancellableContinuation.value ?: return false
+        if (value is CancellableContinuationImpl<*>) return value === requester
+        return true
+    }
 
     /**
      * Claims the continuation for [suspendAtomicCancellableCoroutineReusable] block,

diff --git a/kotlinx-coroutines-core/jvm/test/ReusableCancellableContinuationTest.kt b/kotlinx-coroutines-core/jvm/test/ReusableCancellableContinuationTest.kt
@@ -192,4 +192,17 @@ class ReusableCancellableContinuationTest : TestBase() {
         FieldWalker.assertReachableCount(0, receiver) { it is CancellableContinuation<*> }
         finish(3)
     }
+
+    @Test
+    fun testReusableAndRegularSuspendCancellableCoroutineMemoryLeak() = runTest {
+        val channel =  produce {
+            repeat(10) {
+                send(Unit)
+            }
+        }
+        for (value in channel) {
+            delay(1)
+        }
+        FieldWalker.assertReachableCount(1, coroutineContext[Job], { it is ChildContinuation })
+    }
 }