Release reusability token when suspendCancellableCoroutineReusable's …

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

@@ -331,12 +331,19 @@ internal suspend inline fun <T> suspendCancellableCoroutineReusable(
     crossinline block: (CancellableContinuation<T>) -> Unit
 ): T = suspendCoroutineUninterceptedOrReturn { uCont ->
     val cancellable = getOrCreateCancellableContinuation(uCont.intercepted())
-    block(cancellable)
+    try {
+        block(cancellable)
+    } catch (e: Throwable) {
+        // Here we catch any unexpected exception from user-supplied block (e.g. invariant violation)
+        // and release claimed continuation in order to leave it in a reasonable state (see #3613)
+        cancellable.releaseClaimedReusableContinuation()
+        throw e
+    }
     cancellable.getResult()
 }
 
 internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>): CancellableContinuationImpl<T> {
-    // If used outside of our dispatcher
+    // If used outside our dispatcher
     if (delegate !is DispatchedContinuation<T>) {
         return CancellableContinuationImpl(delegate, MODE_CANCELLABLE)
     }

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

@@ -337,8 +337,8 @@ internal open class CancellableContinuationImpl<in T>(
      * Tries to release reusable continuation. It can fail is there was an asynchronous cancellation,
      * in which case it detaches from the parent and cancels this continuation.
      */
-    private fun releaseClaimedReusableContinuation() {
-        // Cannot be casted if e.g. invoked from `installParentHandleReusable` for context without dispatchers, but with Job in it
+    internal fun releaseClaimedReusableContinuation() {
+        // Cannot be cast if e.g. invoked from `installParentHandleReusable` for context without dispatchers, but with Job in it
         val cancellationCause = (delegate as? DispatchedContinuation<*>)?.tryReleaseClaimedContinuation(this) ?: return
         detachChild()
         cancel(cancellationCause)

kotlinx-coroutines-core/common/src/channels/BufferedChannel.kt

@@ -245,10 +245,8 @@ internal open class BufferedChannel<E>(
      * In case of coroutine cancellation, the element may be undelivered --
      * the [onUndeliveredElement] feature is unsupported in this implementation.
      *
-     * Note that this implementation always invokes [suspendCancellableCoroutineReusable],
-     * as we do not care about broadcasts performance -- they are already deprecated.
      */
-    internal open suspend fun sendBroadcast(element: E): Boolean = suspendCancellableCoroutineReusable { cont ->
+    internal open suspend fun sendBroadcast(element: E): Boolean = suspendCancellableCoroutine { cont ->
         check(onUndeliveredElement == null) {
             "the `onUndeliveredElement` feature is unsupported for `sendBroadcast(e)`"
         }

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

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import org.junit.Test
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicReference
+import kotlin.coroutines.*
+
+// Stresses scenario from #3613
+class ReusableCancellableContinuationInvariantStressTest : TestBase() {
+
+    // Tests have a timeout 10 sec because the bug they catch leads to an infinite spin-loop
+
+    @Test(timeout = 10_000)
+    fun testExceptionFromSuspendReusable() = doTest { /* nothing */ }
+
+
+    @Test(timeout = 10_000)
+    fun testExceptionFromCancelledSuspendReusable() = doTest { it.cancel() }
+
+
+    @Suppress("SuspendFunctionOnCoroutineScope")
+    private inline fun doTest(crossinline block: (Job) -> Unit) {
+        runTest {
+            repeat(10_000) {
+                val latch = CountDownLatch(1)
+                val continuationToResume = AtomicReference<Continuation<Unit>?>(null)
+                val j1 = launch(Dispatchers.Default) {
+                    latch.await()
+                    suspendCancellableCoroutineReusable {
+                        continuationToResume.set(it)
+                        block(coroutineContext.job)
+                        throw CancellationException() // Don't let getResult() chance to execute
+                    }
+                }
+
+                val j2 = launch(Dispatchers.Default) {
+                    latch.await()
+                    while (continuationToResume.get() == null) {
+                        // spin
+                    }
+                    continuationToResume.get()!!.resume(Unit)
+                }
+
+                latch.countDown()
+                joinAll(j1, j2)
+            }
+        }
+    }
+}