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,15 @@ 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 === REUSABLE_CLAIMED -> {
+                    if (_reusableCancellableContinuation.compareAndSet(REUSABLE_CLAIMED, continuation)) return null
+                }
                 state === null -> 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/FieldWalker.kt

@@ -7,6 +7,7 @@ package kotlinx.coroutines
 import java.lang.reflect.*
 import java.util.*
 import java.util.Collections.*
+import kotlin.collections.ArrayList
 
 object FieldWalker {
 
@@ -22,12 +23,6 @@ object FieldWalker {
             val element = stack.removeLast()
             val type = element.javaClass
             type.visit(element, result, stack)
-
-            var superclass = type.superclass
-            while (superclass != Any::class.java && superclass != null) {
-                superclass.visit(element, result, stack)
-                superclass = superclass.superclass
-            }
         }
         return result
     }
@@ -56,9 +51,65 @@ object FieldWalker {
         }
     }
 
-    private fun Class<*>.fields() = declaredFields.filter {
-        !it.type.isPrimitive
-                && !Modifier.isStatic(it.modifiers)
-                && !(it.type.isArray && it.type.componentType.isPrimitive)
+    private fun Class<*>.fields(): List<Field> {
+        val result = ArrayList<Field>()
+        var type = this
+        while (type != Any::class.java) {
+            val fields = type.declaredFields.filter {
+                !it.type.isPrimitive
+                        && !Modifier.isStatic(it.modifiers)
+                        && !(it.type.isArray && it.type.componentType.isPrimitive)
+            }
+            result.addAll(fields)
+            type = type.superclass
+        }
+
+        return result
+    }
+
+    // Debugging-only
+    @Suppress("UNUSED")
+    fun printPath(from: Any, to: Any) {
+        val pathNodes = ArrayList<String>()
+        val visited = newSetFromMap<Any>(IdentityHashMap())
+        visited.add(from)
+        if (findPath(from, to, visited, pathNodes)) {
+            pathNodes.reverse()
+            println(pathNodes.joinToString(" -> ", from.javaClass.simpleName + " -> ", "-> " + to.javaClass.simpleName))
+        } else {
+            println("Path from $from to $to not found")
+        }
+    }
+
+    private fun findPath(from: Any, to: Any, visited: MutableSet<Any>, pathNodes: MutableList<String>): Boolean {
+        if (from === to) {
+            return true
+        }
+
+        val type = from.javaClass
+        if (type.isArray) {
+            if (type.componentType.isPrimitive) return false
+            val array = from as Array<Any?>
+            array.filterNotNull().forEach {
+                if (findPath(it, to, visited, pathNodes)) {
+                    return true
+                }
+            }
+            return false
+        }
+
+        val fields = type.fields()
+        fields.forEach {
+            it.isAccessible = true
+            val value = it.get(from) ?: return@forEach
+            if (!visited.add(value)) return@forEach
+            val found = findPath(value, to, visited, pathNodes)
+            if (found) {
+                pathNodes += from.javaClass.simpleName + ":" + it.name
+                return true
+            }
+        }
+
+        return false
     }
 }

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()
+    }
+}