Flow as Flux transform

Author: SokolovaMaria

reactive/kotlinx-coroutines-reactive/src/flow/FlowAsPublisher.kt

@@ -23,81 +23,83 @@ public fun <T : Any> Flow<T>.asPublisher(): Publisher<T> = FlowAsPublisher(this)
  */
 @Suppress("PublisherImplementation")
 private class FlowAsPublisher<T : Any>(private val flow: Flow<T>) : Publisher<T> {
-
     override fun subscribe(subscriber: Subscriber<in T>?) {
         if (subscriber == null) throw NullPointerException()
         subscriber.onSubscribe(FlowSubscription(flow, subscriber))
     }
+}
 
-    private class FlowSubscription<T>(val flow: Flow<T>, val subscriber: Subscriber<in T>) : Subscription {
-        @Volatile
-        internal var canceled: Boolean = false
-        private val requested = AtomicLong(0L)
-        private val producer: AtomicReference<CancellableContinuation<Unit>?> = AtomicReference()
-
-        // This is actually optimizable
-        private var job = GlobalScope.launch(Dispatchers.Unconfined, start = CoroutineStart.LAZY) {
-            try {
-                consumeFlow()
-                subscriber.onComplete()
-            } catch (e: Throwable) {
-                // Failed with real exception, not due to cancellation
-                if (!coroutineContext[Job]!!.isCancelled) {
-                    subscriber.onError(e)
-                }
+/** @suppress */
+@InternalCoroutinesApi
+public class FlowSubscription<T>(val flow: Flow<T>, val subscriber: Subscriber<in T>) : Subscription {
+    @Volatile
+    internal var canceled: Boolean = false
+    private val requested = AtomicLong(0L)
+    private val producer: AtomicReference<CancellableContinuation<Unit>?> = AtomicReference()
+
+    // This is actually optimizable
+    private var job = GlobalScope.launch(Dispatchers.Unconfined, start = CoroutineStart.LAZY) {
+        try {
+            consumeFlow()
+            subscriber.onComplete()
+        } catch (e: Throwable) {
+            // Failed with real exception, not due to cancellation
+            if (!coroutineContext[Job]!!.isCancelled) {
+                subscriber.onError(e)
             }
         }
+    }
 
-        private suspend fun consumeFlow() {
-            flow.collect { value ->
-                if (!coroutineContext.isActive) {
-                    subscriber.onComplete()
-                    coroutineContext.ensureActive()
-                }
+    private suspend fun consumeFlow() {
+        val ctx = coroutineContext
+        flow.collect { value ->
+            if (!coroutineContext.isActive) {
+                subscriber.onComplete()
+                coroutineContext.ensureActive()
+            }
 
-                if (requested.get() == 0L) {
-                    suspendCancellableCoroutine<Unit> {
-                        producer.set(it)
-                        if (requested.get() != 0L) it.resumeSafely()
-                    }
+            if (requested.get() == 0L) {
+                suspendCancellableCoroutine<Unit> {
+                    producer.set(it)
+                    if (requested.get() != 0L) it.resumeSafely()
                 }
-
-                requested.decrementAndGet()
-                subscriber.onNext(value)
             }
-        }
 
-        override fun cancel() {
-            canceled = true
-            job.cancel()
+            requested.decrementAndGet()
+            subscriber.onNext(value)
         }
+    }
 
-        override fun request(n: Long) {
-            if (n <= 0) {
-                return
-            }
+    override fun cancel() {
+        canceled = true
+        job.cancel()
+    }
+
+    override fun request(n: Long) {
+        if (n <= 0) {
+            return
+        }
 
-            if (canceled) return
+        if (canceled) return
 
-            job.start()
-            var snapshot: Long
-            var newValue: Long
-            do {
-                snapshot = requested.get()
-                newValue = snapshot + n
-                if (newValue <= 0L) newValue = Long.MAX_VALUE
-            } while (!requested.compareAndSet(snapshot, newValue))
+        job.start()
+        var snapshot: Long
+        var newValue: Long
+        do {
+            snapshot = requested.get()
+            newValue = snapshot + n
+            if (newValue <= 0L) newValue = Long.MAX_VALUE
+        } while (!requested.compareAndSet(snapshot, newValue))
 
-            val prev = producer.get()
-            if (prev == null || !producer.compareAndSet(prev, null)) return
-            prev.resumeSafely()
-        }
+        val prev = producer.get()
+        if (prev == null || !producer.compareAndSet(prev, null)) return
+        prev.resumeSafely()
+    }
 
-        private fun CancellableContinuation<Unit>.resumeSafely() {
-            val token = tryResume(Unit)
-            if (token != null) {
-                completeResume(token)
-            }
+    private fun CancellableContinuation<Unit>.resumeSafely() {
+        val token = tryResume(Unit)
+        if (token != null) {
+            completeResume(token)
         }
     }
-}
+}

reactive/kotlinx-coroutines-reactor/src/FlowAsFlux.kt

@@ -0,0 +1,21 @@
+package kotlinx.coroutines.reactor
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.reactive.flow.FlowSubscription
+import reactor.core.CoreSubscriber
+import reactor.core.publisher.Flux
+
+/**
+ * Transforms the given flow to [Flux].
+ */
+@ExperimentalCoroutinesApi
+public fun <T: Any> Flow<T>.asFlux(): Flux<T> = FlowAsFlux(this)
+
+private class FlowAsFlux<T : Any>(private val flow: Flow<T>) : Flux<T>() {
+    override fun subscribe(subscriber: CoreSubscriber<in T>?) {
+        if (subscriber == null) throw NullPointerException()
+        subscriber.onSubscribe(FlowSubscription(flow.flowOn(subscriber.currentContext().asCoroutineContext()), subscriber))
+    }
+}

reactive/kotlinx-coroutines-reactor/src/Flux.kt

@@ -74,4 +74,4 @@ private fun <T> reactorPublish(
     val coroutine = PublisherCoroutine(newContext, subscriber)
     subscriber.onSubscribe(coroutine) // do it first (before starting coroutine), to avoid unnecessary suspensions
     coroutine.start(CoroutineStart.DEFAULT, coroutine, block)
-}
+}

reactive/kotlinx-coroutines-reactor/test/FlowAsFluxTest.kt

@@ -0,0 +1,27 @@
+package kotlinx.coroutines.reactor
+
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.reactive.*
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import reactor.core.publisher.Mono
+import reactor.util.context.Context
+import kotlin.test.assertEquals
+
+class FlowAsFluxTest {
+    @Test
+    fun testFlowToFluxContextPropagation() = runBlocking<Unit> {
+        val flux = flow<String> {
+            (1..4).forEach { i -> emit(m(i).awaitFirst()) }
+        }   .asFlux()
+            .subscriberContext(Context.of(1, "1"))
+            .subscriberContext(Context.of(2, "2", 3, "3", 4, "4"))
+        var i = 0
+        flux.subscribe { str -> i++; println(str); assertEquals(str, i.toString()) }
+    }
+
+    private fun m(i: Int): Mono<String> = mono {
+        val ctx = coroutineContext[ReactorContext]?.context
+        ctx?.getOrDefault(i, "noValue")
+    }
+}