Fix subject cancellation, add FlatMapDrop

Author: akarnokd

README.md

@@ -11,7 +11,7 @@ Extensions to the Kotlin Flow library.
 
 ```groovy
 dependencies {
-    implementation "com.github.akarnokd:kotlin-flow-extensions:0.0.2"
+    implementation "com.github.akarnokd:kotlin-flow-extensions:0.0.3"
 }
 ```
 
@@ -34,6 +34,8 @@ Table of contents
   - `Flow.replay`
   - `Flow.startCollectOn`
   - `Flow.takeUntil`
+  - `Flow.onBackpressureDrop`
+  - [`Flow.flatMapDrop`](#flowflatmapdrop)
 - `ParallelFlow` operators (`FlowExtensions`)
   - `ParallelFlow.concatMap`
   - `ParallelFlow.filter`
@@ -150,3 +152,33 @@ runBlocking {
     job.join()
 }
 ```
+
+## Flow.flatMapDrop
+
+Maps the upstream value into a `Flow` and relays its items while ignoring further upstream items until the current
+inner `Flow` completes.
+
+```kotlin
+import hu.akarnokd.kotlin.flow.*
+
+range(1, 10)
+.map {
+    delay(100)
+    it
+}
+.flatMapDrop {
+    range(it * 100, 5)
+            .map {
+                delay(30)
+                it
+            }
+}
+.assertResult(
+        100, 101, 102, 103, 104,
+        300, 301, 302, 303, 304,
+        500, 501, 502, 503, 504,
+        700, 701, 702, 703, 704,
+        900, 901, 902, 903, 904
+)
+```
+

gradle.properties

@@ -1,5 +1,5 @@
 GROUP=com.github.akarnokd
-VERSION_NAME=0.0.2
+VERSION_NAME=0.0.3
 
 POM_ARTIFACT_ID=kotlin-flow-extensions
 POM_NAME=Kotlin Flow Extensions

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

@@ -16,10 +16,13 @@
 
 package hu.akarnokd.kotlin.flow
 
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.AbstractFlow
 import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.isActive
 import java.util.concurrent.atomic.AtomicReference
+import kotlin.coroutines.coroutineContext
 
 /**
  * Caches one item and replays it to fresh collectors.
@@ -134,7 +137,11 @@ class BehaviorSubject<T> : AbstractFlow<T>, SubjectAPI<T> {
 
             if (curr.value != NONE) {
                 try {
-                    collector.emit(curr.value)
+                    if (coroutineContext.isActive) {
+                        collector.emit(curr.value)
+                    } else {
+                        throw CancellationException()
+                    }
                 } catch (exc: Throwable) {
                     remove(inner)
 
@@ -160,7 +167,11 @@ class BehaviorSubject<T> : AbstractFlow<T>, SubjectAPI<T> {
                 }
 
                 try {
-                    collector.emit(next.value)
+                    if (coroutineContext.isActive) {
+                        collector.emit(next.value)
+                    } else {
+                        throw CancellationException()
+                    }
                 } catch (exc: Throwable) {
                     remove(inner)
 

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

@@ -164,6 +164,13 @@ fun <T> Flow<T>.toList() : Flow<List<T>> {
 @FlowPreview
 fun <T> Flow<T>.onBackpressurureDrop() : Flow<T> = FlowOnBackpressureDrop(this)
 
+/**
+ * Maps items from the upstream to [Flow] and relays its items while dropping upstream items
+ * until the current inner [Flow] completes.
+ */
+@FlowPreview
+fun <T, R> Flow<T>.flatMapDrop(mapper: suspend (T) -> Flow<R>) : Flow<R> = FlowFlatMapDrop(this, mapper)
+
 // -----------------------------------------------------------------------------------------
 // Parallel Extensions
 // -----------------------------------------------------------------------------------------

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

@@ -16,11 +16,14 @@
 
 package hu.akarnokd.kotlin.flow
 
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.AbstractFlow
 import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.isActive
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
+import kotlin.coroutines.coroutineContext
 
 /**
  * Caches and replays some or all items to collectors.
@@ -240,8 +243,12 @@ class ReplaySubject<T> : AbstractFlow<T>, SubjectAPI<T> {
 
                 if (!empty) {
                     try {
-                        consumer.consumer.emit(list[consumer.index.toInt()])
-                        consumer.index++
+                        if (coroutineContext.isActive) {
+                            consumer.consumer.emit(list[consumer.index.toInt()])
+                            consumer.index++
+                        } else {
+                            throw CancellationException()
+                        }
                     } catch (ex: Throwable) {
                         consumer.parent.remove(consumer)
 
@@ -317,8 +324,12 @@ class ReplaySubject<T> : AbstractFlow<T>, SubjectAPI<T> {
 
                 if (!empty) {
                     try {
-                        consumer.consumer.emit(next.value!!)
-                        consumer.node = next
+                        if (coroutineContext.isActive) {
+                            consumer.consumer.emit(next.value!!)
+                            consumer.node = next
+                        } else {
+                            throw CancellationException()
+                        }
                     } catch (ex: Throwable) {
                         consumer.parent.remove(consumer)
 
@@ -435,8 +446,12 @@ class ReplaySubject<T> : AbstractFlow<T>, SubjectAPI<T> {
 
                 if (!empty) {
                     try {
-                        consumer.consumer.emit(next.value!!)
-                        consumer.node = next
+                        if (coroutineContext.isActive) {
+                            consumer.consumer.emit(next.value!!)
+                            consumer.node = next
+                        } else {
+                            throw CancellationException()
+                        }
                     } catch (ex: Throwable) {
                         consumer.parent.remove(consumer)
 

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

@@ -86,7 +86,6 @@ open class ResumableCollector<T> : Resumable() {
                     if (coroutineContext.isActive) {
                         collector.emit(v)
                     } else {
-                        onComplete?.invoke(this)
                         throw CancellationException()
                     }
                 } catch (exc: Throwable) {

src/main/kotlin/hu/akarnokd/kotlin/flow/impl/FlowFlatMapDrop.kt

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2019 David Karnok
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package hu.akarnokd.kotlin.flow.impl
+
+import hu.akarnokd.kotlin.flow.Resumable
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.AbstractFlow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.flow.collect
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicReference
+
+@FlowPreview
+internal class FlowFlatMapDrop<T, R>(private val source: Flow<T>, private val mapper: suspend (T) -> Flow<R>) : AbstractFlow<R>() {
+    @ExperimentalCoroutinesApi
+    @InternalCoroutinesApi
+    override suspend fun collectSafely(collector: FlowCollector<R>) {
+        coroutineScope {
+
+            val resume = Resumable();
+            val consumerReady = AtomicBoolean(true)
+            val value = AtomicReference<T>()
+            val hasValue = AtomicBoolean()
+            val done = AtomicBoolean()
+            val error = AtomicReference<Throwable>()
+
+            val job = launch {
+                try {
+                    source.collect {
+                        if (consumerReady.get()) {
+                            consumerReady.set(false)
+                            value.lazySet(it)
+                            hasValue.set(true);
+                            resume.resume()
+                        }
+                    }
+                    done.set(true)
+                } catch (ex: Throwable) {
+                    error.set(ex)
+                }
+                resume.resume()
+            }
+
+            while (coroutineContext.isActive) {
+                resume.await()
+
+                val d = done.get()
+                val e = error.get()
+                val h = hasValue.get()
+                val v = value.get()
+
+                if (e != null && !h) {
+                    throw e;
+                }
+
+                if (d && !h) {
+                    break;
+                }
+
+                if (h) {
+                    value.lazySet(null)
+                    hasValue.set(false)
+                    try {
+                        mapper(v).collect {
+                            collector.emit(it)
+                        }
+                    } catch (ex: Throwable) {
+                        job.cancel()
+                        throw ex
+                    }
+                }
+
+                consumerReady.set(true);
+            }
+        }
+    }
+}

src/main/kotlin/hu/akarnokd/kotlin/flow/impl/FlowOnBackpressureDrop.kt

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2019 David Karnok
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package hu.akarnokd.kotlin.flow.impl
 
 import hu.akarnokd.kotlin.flow.Resumable

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

@@ -363,4 +363,88 @@ class BehaviorSubjectTest {
             assertEquals<Any?>(listOf(null), result)
         }
     }
+
+    @Test
+    fun cancelledConsumer() = runBlocking {
+        withSingle {
+            val subject = BehaviorSubject<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)
+            }
+
+            // wait for the subject to finish
+            for (i in 1..1000) {
+                if (job1.isCancelled) {
+                    break;
+                }
+                delay(10)
+            }
+
+            assertEquals(true, job1.isCancelled)
+            assertEquals(expected, counter1.get())
+            assertEquals(0, subject.collectorCount())
+        }
+
+    }
+
+    @Test
+    fun cancelledOneCollectorSecondCompletes() = runBlocking {
+        withSingle {
+            val subject = BehaviorSubject<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())
+        }
+
+    }
 }

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

@@ -301,6 +301,14 @@ class PublishSubjectTest {
                 subject.emit(i)
             }
 
+            // wait for the subject to finish
+            for (i in 1..1000) {
+                if (job1.isCancelled) {
+                    break;
+                }
+                delay(10)
+            }
+
             assertEquals(true, job1.isCancelled)
             assertEquals(expected, counter1.get())
             assertEquals(0, subject.collectorCount())