From 6e66695bbbfc872c38f8e9e88d1a9d00cb43169f Mon Sep 17 00:00:00 2001
From: Roman Elizarov <elizarov@gmail.com>
Date: Tue, 31 Mar 2020 16:10:12 +0300
Subject: [PATCH] Fix leak of subscription reference in closed
 ArrayBroadcastChannel (#1885)

When ArrayBroadcastChannel was closed it was still retaining a reference to its subscription (even if that subscription was cancelled) while, in fact, either closing a broadcast channel or cancelling subscription should remove the reference.

This is no problem with ConflatedBroadcastChannel, but it is added to the test for completeness.
---
 .../src/channels/ArrayBroadcastChannel.kt     |  4 ++-
 .../test/channels/BroadcastChannelLeakTest.kt | 34 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 kotlinx-coroutines-core/jvm/test/channels/BroadcastChannelLeakTest.kt

diff --git a/kotlinx-coroutines-core/common/src/channels/ArrayBroadcastChannel.kt b/kotlinx-coroutines-core/common/src/channels/ArrayBroadcastChannel.kt
index 19334ea706..155652fd6f 100644
--- a/kotlinx-coroutines-core/common/src/channels/ArrayBroadcastChannel.kt
+++ b/kotlinx-coroutines-core/common/src/channels/ArrayBroadcastChannel.kt
@@ -218,13 +218,15 @@ internal class ArrayBroadcastChannel<E>(
         override val isBufferAlwaysFull: Boolean get() = error("Should not be used")
         override val isBufferFull: Boolean get() = error("Should not be used")
 
-        override fun onCancelIdempotent(wasClosed: Boolean) {
+        override fun close(cause: Throwable?): Boolean {
+            val wasClosed = super.close(cause)
             if (wasClosed) {
                 broadcastChannel.updateHead(removeSub = this)
                 subLock.withLock {
                     subHead = broadcastChannel.tail
                 }
             }
+            return wasClosed
         }
 
         // returns true if subHead was updated and broadcast channel's head must be checked
diff --git a/kotlinx-coroutines-core/jvm/test/channels/BroadcastChannelLeakTest.kt b/kotlinx-coroutines-core/jvm/test/channels/BroadcastChannelLeakTest.kt
new file mode 100644
index 0000000000..66b08c74e4
--- /dev/null
+++ b/kotlinx-coroutines-core/jvm/test/channels/BroadcastChannelLeakTest.kt
@@ -0,0 +1,34 @@
+package kotlinx.coroutines.channels
+
+import kotlinx.coroutines.*
+import org.junit.Test
+import kotlin.test.*
+
+class BroadcastChannelLeakTest : TestBase() {
+    @Test
+    fun testArrayBroadcastChannelSubscriptionLeak() {
+        checkLeak { ArrayBroadcastChannel(1) }
+    }
+
+    @Test
+    fun testConflatedBroadcastChannelSubscriptionLeak() {
+        checkLeak { ConflatedBroadcastChannel() }
+    }
+
+    enum class TestKind { BROADCAST_CLOSE, SUB_CANCEL, BOTH }
+
+    private fun checkLeak(factory: () -> BroadcastChannel<String>) = runTest {
+        for (kind in TestKind.values()) {
+            val broadcast = factory()
+            val sub = broadcast.openSubscription()
+            broadcast.send("OK")
+            assertEquals("OK", sub.receive())
+            // now close broadcast
+            if (kind != TestKind.SUB_CANCEL) broadcast.close()
+            // and then cancel subscription
+            if (kind != TestKind.BROADCAST_CLOSE) sub.cancel()
+            // subscription should not be reachable from the channel anymore
+            FieldWalker.assertReachableCount(0, broadcast) { it === sub }
+        }
+    }
+}
\ No newline at end of file