Ensure subject removes cancelled consumer (#17)

* Ensure subject removes cancelled consumer

Fixes #11

* fixing formatting

Author: KVelasco

src/main/kotlin/hu/akarnokd/kotlin/flow/ResumableCollector.kt

@@ -16,7 +16,10 @@
 
 package hu.akarnokd.kotlin.flow
 
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.isActive
+import kotlin.coroutines.coroutineContext
 
 /**
  * A collector that hosts a signal (value/error/completion)
@@ -66,7 +69,7 @@ open class ResumableCollector<T> : Resumable() {
         consumerReady.resume()
     }
 
-    suspend fun drain(collector: FlowCollector<T>, onCrash: ((ResumableCollector<T>) -> Unit)? = null) {
+    suspend fun drain(collector: FlowCollector<T>, onComplete: ((ResumableCollector<T>) -> Unit)? = null) {
         while (true) {
 
             readyConsumer()
@@ -80,9 +83,14 @@ open class ResumableCollector<T> : Resumable() {
                 hasValue = false
 
                 try {
-                    collector.emit(v)
+                    if (coroutineContext.isActive) {
+                        collector.emit(v)
+                    } else {
+                        onComplete?.invoke(this)
+                        throw CancellationException()
+                    }
                 } catch (exc: Throwable) {
-                    onCrash?.invoke(this)
+                    onComplete?.invoke(this)
 
                     readyConsumer() // unblock waiters
                     throw exc

src/test/kotlin/hu/akarnokd/kotlin/flow/PublishSubjectTest.kt

@@ -274,4 +274,80 @@ class PublishSubjectTest {
 
         assertEquals(0, counter1.get())
     }
+
+    @Test
+    fun cancelledConsumer() = runBlocking {
+        withSingle {
+            val subject = PublishSubject<Int>()
+
+            val expected = 3
+            val n = 10
+
+            val counter1 = AtomicInteger()
+
+            val job1 = launch(it.asCoroutineDispatcher()) {
+                subject.collect {
+                    if (counter1.incrementAndGet() == expected) {
+                        cancel()
+                    }
+                }
+            }
+
+            while (!subject.hasCollectors()) {
+                delay(1)
+            }
+
+            for (i in 1..n) {
+                subject.emit(i)
+            }
+
+            assertEquals(true, job1.isCancelled)
+            assertEquals(expected, counter1.get())
+            assertEquals(0, subject.collectorCount())
+        }
+
+    }
+
+    @Test
+    fun cancelledOneCollectorSecondCompletes() = runBlocking {
+        withSingle {
+            val subject = PublishSubject<Int>()
+
+            val expected = 3
+            val n = 10
+
+            val counter1 = AtomicInteger()
+            val counter2 = AtomicInteger()
+
+            val job1 = launch(it.asCoroutineDispatcher()) {
+                subject.collect {
+                    if (counter1.incrementAndGet() == expected) {
+                        cancel()
+                    }
+                }
+            }
+
+            val job2 = launch(it.asCoroutineDispatcher()) {
+                subject.collect { counter2.incrementAndGet() }
+            }
+
+            while (subject.collectorCount() != 2) {
+                delay(1)
+            }
+
+            for (i in 1..n) {
+                subject.emit(i)
+            }
+
+            subject.complete()
+            job2.join()
+
+            assertEquals(true, job1.isCancelled)
+            assertEquals(true, job2.isCompleted)
+            assertEquals(expected, counter1.get())
+            assertEquals(n, counter2.get())
+            assertEquals(0, subject.collectorCount())
+        }
+
+    }
 }