Ensure ConsumeAsFlow does not retain reference to the last element of the flow with test

Author: qwwdfsad

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

@@ -31,7 +31,7 @@ internal class DispatchedContinuation<in T>(
      * Possible states of reusability:
      *
      * 1) `null`. Cancellable continuation wasn't yet attempted to be reused or
-     *     way used and then invalidated (e.g. because of the cancellation).
+     *     was 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:
      *    ```
@@ -91,7 +91,7 @@ internal class DispatchedContinuation<in T>(
                         return state as CancellableContinuationImpl<T>
                     }
                 }
-                else ->  error("Inconsistent state $state")
+                else -> error("Inconsistent state $state")
             }
         }
     }
@@ -114,13 +114,14 @@ internal class DispatchedContinuation<in T>(
         _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 === REUSABLE_CLAIMED -> {
+                    if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, continuation)) return null
+                }
                 state is Throwable -> {
                     require(_reusableCancellableContinuation.compareAndSet(state, null))
                     return state
                 }
-                else -> return null // Is not reusable
+                else -> error("Inconsistent state $state")
             }
         }
     }

kotlinx-coroutines-core/jvm/test/flow/ConsumeAsFlowLeakTest.kt

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import org.junit.Test
+import kotlin.test.*
+
+class ConsumeAsFlowLeakTest : TestBase() {
+
+    private data class Box(val i: Int)
+
+    // In companion to avoid references through runTest
+    companion object {
+        private val first = Box(4)
+        private val second = Box(5)
+    }
+
+    // @Test //ignored until KT-33986
+    fun testReferenceIsNotRetained() = testReferenceNotRetained(true)
+
+    @Test
+    fun testReferenceIsNotRetainedNoSuspension() = testReferenceNotRetained(false)
+
+    private fun testReferenceNotRetained(shouldSuspendOnSend: Boolean) = runTest {
+        val channel = BroadcastChannel<Box>(1)
+        val job = launch {
+            expect(2)
+            channel.asFlow().collect {
+                expect(it.i)
+            }
+        }
+
+        expect(1)
+        yield()
+        expect(3)
+        channel.send(first)
+        if (shouldSuspendOnSend) yield()
+        channel.send(second)
+        yield()
+        assertEquals(0, FieldWalker.walk(channel).count { it === second })
+        finish(6)
+        job.cancelAndJoin()
+    }
+}