Publish(expectedCollectors) overload

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.5"
+    implementation "com.github.akarnokd:kotlin-flow-extensions:0.0.6"
 }
 ```
 
@@ -30,7 +30,7 @@ Table of contents
   - `Flow.concatWith`
   - `Flow.groupBy`
   - `Flow.parallel`
-  - `Flow.publish`
+  - [`Flow.publish`](#flowpublish)
   - `Flow.replay`
   - `Flow.startCollectOn`
   - `Flow.takeUntil`
@@ -182,3 +182,30 @@ range(1, 10)
 )
 ```
 
+## Flow.publish
+
+Shares a single connection to the upstream source which can be consumed by many collectors inside a `transform` function,
+which then yields the resulting items for the downstream.
+
+Effectively, one collector to the output `Flow<R>` will trigger exactly one collection of the upstream `Flow<T>`. Inside
+the `transformer` function though, the presented `Flow<T>` can be collected as many times as needed; it won't trigger
+new collections towards the upstream but share items to all inner collectors as they become available.
+
+Unfortunately, the suspending nature of coroutines/`Flow` doesn't give a clear indication when the `transformer` chain
+has been properly established, which can result in item loss or run-to-completion without any item being collected.
+If the number of the inner collectors inside `transformer` can be known, the `publish(expectedCollectors)` overload
+can be used to hold back the upstream until the expected number of collectors have started/ready collecting items.
+
+#### Example:
+
+```kotlin
+    range(1, 5)
+    .publish(2) { 
+         shared -> merge(shared.filter { it % 2 == 0 }, shared.filter { it % 2 != 0 }) 
+    }
+    .assertResult(1, 2, 3, 4, 5)
+```
+
+In the example, it is known `merge` will establish 2 collectors, thus the `publish` can be instructed to await those 2.
+Without the argument, `range` would rush through its items as `merge` doesn't start collecting in time, causing an
+empty result list.

gradle.properties

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

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

@@ -28,10 +28,35 @@ import java.util.concurrent.TimeUnit
  * Shares a single collector towards the upstream source and multicasts
  * values to any number of consumers which then can produce the output
  * flow of values.
+ *
+ * Note that due to how coroutines/[Flow] are implemented, it is not guaranteed
+ * the [transform] function connects the upstream with the downstream in time,
+ * causing item loss or even run-to-completion without any single upstream item
+ * being collected and transformed. To avoid such scenarios, use the
+ * `publish(expectedCollectors)` overload.
  */
 @FlowPreview
 fun <T, R> Flow<T>.publish(transform: suspend (Flow<T>) -> Flow<R>) : Flow<R> =
-    FlowMulticastFunction(this, { /* MulticastSubject()*/ PublishSubject() }, transform)
+    FlowMulticastFunction(this, { PublishSubject() }, transform)
+
+/**
+ * Shares a single collector towards the upstream source and multicasts
+ * values to any number of consumers which then can produce the output
+ * flow of values.
+ *
+ * Note that due to how coroutines/[Flow] are implemented, it is not guaranteed
+ * the [transform] function connects the upstream with the downstream in time,
+ * causing item loss or even run-to-completion without any single upstream item
+ * being collected and transformed. To avoid such scenarios, specify the
+ * [expectedCollectors] to delay the collection of the upstream until the number
+ * of inner collectors has reached the specified number.
+ *
+ * @param expectedCollectors the number of collectors to wait for before resuming the source, allowing
+ * the desired number of collectors to arrive and be ready for the upstream items
+ */
+@FlowPreview
+fun <T, R> Flow<T>.publish(expectedCollectors: Int, transform: suspend (Flow<T>) -> Flow<R>) : Flow<R> =
+        FlowMulticastFunction(this, { /* MulticastSubject()*/ MulticastSubject(expectedCollectors) }, transform)
 
 /**
  * Shares a single collector towards the upstream source and multicasts

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

@@ -30,178 +30,137 @@ import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.atomic.AtomicReference
 
 /**
- * A subject implementation that dispatches signals to multiple
- * consumers or buffers them until such consumers arrive.
+ * A subject implementation that awaits a certain number of collectors
+ * to start consuming, then allows the producer side to deliver items
+ * to them.
  *
  * @param <T> the element type of the [Flow]
  * @param bufferSize the number of items to buffer until consumers arrive
  */
 @FlowPreview
-class MulticastSubject<T>(private val bufferSize: Int = 32) : AbstractFlow<T>(), SubjectAPI<T> {
+class MulticastSubject<T>(private val expectedCollectors: Int) : AbstractFlow<T>(), SubjectAPI<T> {
 
-    val queue = ConcurrentLinkedQueue<T>()
+    val collectors = AtomicReference<Array<ResumableCollector<T>>>(EMPTY as Array<ResumableCollector<T>>)
 
-    val availableQueue = AtomicInteger(bufferSize)
+    val producer = Resumable()
 
-    val consumers = AtomicReference(EMPTY as Array<ResumableCollector<T>>)
-
-    val producerAwait = Resumable()
-
-    val wip = AtomicInteger()
+    val remainingCollectors = AtomicInteger(expectedCollectors)
 
     @Volatile
-    var error: Throwable? = null
+    var terminated : Throwable? = null
 
     override suspend fun emit(value: T) {
-        while (availableQueue.get() == 0) {
-            producerAwait.await()
+        awaitCollectors()
+        for (collector in collectors.get()) {
+            try {
+                collector.next(value)
+            } catch (ex: CancellationException) {
+                remove(collector)
+            }
         }
-        queue.offer(value)
-        availableQueue.decrementAndGet();
-        drain()
     }
 
     override suspend fun emitError(ex: Throwable) {
-        error = ex
-        drain()
+        //awaitCollectors()
+        terminated = ex;
+        for (collector in collectors.getAndSet(TERMINATED as Array<ResumableCollector<T>>)) {
+            try {
+                collector.error(ex)
+            } catch (_: CancellationException) {
+                // ignored at this point
+            }
+        }
     }
 
     override suspend fun complete() {
-        error = DONE
-        drain()
+        //awaitCollectors()
+        terminated = DONE
+        for (collector in collectors.getAndSet(TERMINATED as Array<ResumableCollector<T>>)) {
+            try {
+                collector.complete()
+            } catch (_: CancellationException) {
+                // ignored at this point
+            }
+        }
     }
 
     override fun hasCollectors(): Boolean {
-        return consumers.get().isNotEmpty();
+        return collectors.get().isNotEmpty()
     }
 
     override fun collectorCount(): Int {
-        return consumers.get().size;
+        return collectors.get().size
     }
 
-    override suspend fun collectSafely(collector: FlowCollector<T>) {
-        val c = ResumableCollector<T>()
-        if (add(c)) {
-            c.readyConsumer()
-            drain()
-            c.drain(collector) { remove(it) }
-        } else {
-            val ex = error
-            if (ex != null && ex != DONE) {
-                throw ex
-            }
+    private suspend fun awaitCollectors() {
+        if (remainingCollectors.get() != 0) {
+            producer.await()
         }
     }
 
-    private fun add(collector: ResumableCollector<T>) : Boolean {
+    @Suppress("UNCHECKED_CAST", "")
+    private fun add(inner: ResumableCollector<T>) : Boolean {
         while (true) {
-            val a = consumers.get()
-            if (a == TERMINATED) {
+
+            val a = collectors.get()
+            if (a as Any == TERMINATED as Any) {
                 return false
             }
-            val b = Array<ResumableCollector<T>>(a.size + 1) { idx ->
-                if (idx < a.size) a[idx] else collector
-            }
-            if (consumers.compareAndSet(a, b)) {
+            val n = a.size
+            val b = a.copyOf(n + 1)
+            b[n] = inner
+            if (collectors.compareAndSet(a, b as Array<ResumableCollector<T>>)) {
                 return true
             }
         }
     }
-    private fun remove(collector: ResumableCollector<T>) {
+
+    @Suppress("UNCHECKED_CAST")
+    private fun remove(inner: ResumableCollector<T>) {
         while (true) {
-            val a = consumers.get()
+            val a = collectors.get()
             val n = a.size
             if (n == 0) {
                 return
             }
-            var j = -1
-            for (i in 0 until n) {
-                if (a[i] == collector) {
-                    j = i
-                    break
-                }
-            }
+
+            val j = a.indexOf(inner)
             if (j < 0) {
-                return;
+                return
             }
+
             var b = EMPTY as Array<ResumableCollector<T>?>
             if (n != 1) {
-                b = Array<ResumableCollector<T>?>(n - 1) { null }
+                b = Array(n - 1) { null }
                 System.arraycopy(a, 0, b, 0, j)
                 System.arraycopy(a, j + 1, b, j, n - j - 1)
             }
-            if (consumers.compareAndSet(a, b as Array<ResumableCollector<T>>)) {
-                return;
-            }
-        }
-    }
-
-    private suspend fun drain() {
-        if (wip.getAndIncrement() != 0) {
-            return;
-        }
-
-        while (true) {
-            val collectors = consumers.get()
-            if (collectors.isNotEmpty()) {
-                val ex = error;
-                val v = queue.poll();
-
-                if (v == null && ex != null) {
-                    finish(ex)
-                }
-                else if (v != null) {
-                    var k = 0;
-                    for (collector in consumers.get()) {
-                        try {
-                            //println("MulticastSubject -> [$k]: $v")
-                            collector.next(v)
-                        } catch (ex: CancellationException) {
-                            remove(collector);
-                        }
-                        k++
-                    }
-                    availableQueue.getAndIncrement()
-                    producerAwait.resume()
-                    continue
-                }
-            } else {
-                val ex = error;
-                if (ex != null && queue.isEmpty()) {
-                    finish(ex)
-                }
-            }
-            if (wip.decrementAndGet() == 0) {
-                break
+            if (collectors.compareAndSet(a, b as Array<ResumableCollector<T>>)) {
+                return
             }
         }
     }
 
-    private suspend fun finish(ex: Throwable) {
-        if (ex == DONE) {
-            for (collector in consumers.getAndSet(TERMINATED as Array<ResumableCollector<T>>)) {
-                try {
-                    collector.complete()
-                } catch (_: CancellationException) {
-                    // ignored
-                }
+    override suspend fun collectSafely(collector: FlowCollector<T>) {
+        val rc = ResumableCollector<T>()
+        if (add(rc)) {
+            if (remainingCollectors.decrementAndGet() == 0) {
+                producer.resume()
             }
+            rc.drain(collector) { remove(it) }
         } else {
-            for (collector in consumers.getAndSet(TERMINATED as Array<ResumableCollector<T>>)) {
-                try {
-                    collector.error(ex)
-                } catch (_: CancellationException) {
-                    // ignored
-                }
+            val ex = terminated;
+            if (ex != null && ex != DONE) {
+                throw ex
             }
         }
     }
 
     companion object {
-        val DONE: Throwable = Throwable("Subject Completed")
+        val EMPTY = arrayOf<ResumableCollector<Any>>()
 
-        val EMPTY = arrayOf<ResumableCollector<Any>>();
+        val TERMINATED = arrayOf<ResumableCollector<Any>>()
 
-        val TERMINATED = arrayOf<ResumableCollector<Any>>();
+        val DONE = Throwable("Subject completed")
     }
 }