Add Flow.amb

Author: akarnokd

README.md

@@ -40,6 +40,7 @@ Table of contents
   - `Flow.onBackpressureDrop`
   - [`Flow.flatMapDrop`](#flowflatmapdrop)
   - [`Flow.concatMapEager`](#flowconcatmapeager)
+  - [`Flow.amb`](#flowamb)
 - `ParallelFlow` operators (`FlowExtensions`)
   - `ParallelFlow.concatMap`
   - `ParallelFlow.filter`
@@ -308,4 +309,17 @@ range(1, 5)
         40, 41, 42, 43, 44,
         50, 51, 52, 53, 54
 )
+```
+
+## Flow.amb
+
+Starts collecting all source [Flow]s and relays the items of the first one to emit an item,
+cancelling the rest.
+
+```kotlin
+amb(
+    range(1, 5).onStart { delay(1000) },
+    range(6, 5).onStart { delay(100) }
+)
+.assertResult(6, 7, 8, 9, 10)
 ```

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

@@ -22,7 +22,9 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.*
+import java.util.*
 import java.util.concurrent.TimeUnit
+import kotlin.collections.ArrayList
 
 /**
  * Shares a single collector towards the upstream source and multicasts
@@ -231,6 +233,22 @@ fun <T> concatArrayEager(vararg sources: Flow<T>) : Flow<T> = FlowConcatArrayEag
 @FlowPreview
 fun <T, R> Flow<T>.concatMapEager(mapper: suspend (T) -> Flow<R>) : Flow<R> = FlowConcatMapEager(this, mapper)
 
+/**
+ * Starts collecting all source [Flow]s and relays the items of the first one to emit an item,
+ * cancelling the rest.
+ * @param sources the [Iterable] sequence of [Flow]s
+ */
+@FlowPreview
+fun <T> amb(sources: Iterable<Flow<T>>) : Flow<T> = FlowAmbIterable(sources)
+
+/**
+ * Starts collecting all source [Flow]s and relays the items of the first one to emit an item,
+ * cancelling the rest.
+ * @param sources the array of [Flow]s
+ */
+@FlowPreview
+fun <T> amb(vararg sources: Flow<T>) : Flow<T> = FlowAmbIterable(listOf(*sources))
+
 // -----------------------------------------------------------------------------------------
 // Parallel Extensions
 // -----------------------------------------------------------------------------------------

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

@@ -0,0 +1,51 @@
+package hu.akarnokd.kotlin.flow.impl
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.flow.collect
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicInteger
+
+@FlowPreview
+class FlowAmbIterable<T>(private val sources: Iterable<Flow<T>>) : Flow<T> {
+    @InternalCoroutinesApi
+    override suspend fun collect(collector: FlowCollector<T>) {
+        val winner = AtomicInteger()
+        val jobs = ConcurrentHashMap<Job, Int>()
+        coroutineScope {
+            var i = 1
+            for (source in sources) {
+                val idx = i
+                val job = launch {
+                    source.collect {
+                        val w = winner.get()
+                        if (w == idx) {
+                            collector.emit(it)
+                        } else if (w == 0 && winner.compareAndSet(0, idx)) {
+                            for (j in jobs.entries) {
+                                if (j.value != idx) {
+                                    j.key.cancel()
+                                }
+                            }
+
+                            collector.emit(it)
+                        } else {
+                            throw CancellationException()
+                        }
+                    }
+                }
+
+                jobs[job] = i
+                val w = winner.get()
+                if (w != 0 && w != i) {
+                    jobs.remove(job)
+                    job.cancel()
+                    break
+                }
+
+                i++
+            }
+        }
+    }
+}

src/test/kotlin/hu/akarnokd/kotlin/flow/impl/FlowAmbIterableTest.kt

@@ -0,0 +1,53 @@
+package hu.akarnokd.kotlin.flow.impl
+
+import hu.akarnokd.kotlin.flow.amb
+import hu.akarnokd.kotlin.flow.assertResult
+import hu.akarnokd.kotlin.flow.concatArrayEager
+import hu.akarnokd.kotlin.flow.range
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.test.assertEquals
+
+@FlowPreview
+class FlowAmbIterableTest {
+    @Test
+    fun basic1() = runBlocking {
+        amb(
+                range(1, 5).onStart { delay(1000) },
+                range(6, 5).onStart { delay(100) }
+        )
+        .assertResult(6, 7, 8, 9, 10)
+    }
+
+    @Test
+    fun basic2() = runBlocking {
+        amb(
+                range(1, 5).onStart { delay(100) },
+                range(6, 5).onStart { delay(1000) }
+        )
+                .assertResult(1, 2, 3, 4, 5)
+    }
+
+    @Test
+    fun basic3() = runBlocking {
+        val counter = AtomicInteger()
+        amb(
+                range(1, 5).onEach { delay(100) },
+                range(6, 5)
+                        .onEach {
+                            delay(200)
+                            counter.getAndIncrement()
+                        }
+
+        )
+                .take(3)
+                .assertResult(1, 2, 3)
+
+        assertEquals(0, counter.get())
+    }
+
+}