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.

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

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