Optimize Publisher.asFlow fusion and implementation size

* Calling Publisher.asFlow().produceIn(...) uses a single channel.
* Implementation fully reuses Publisher.openSubscription code.

Author: elizarov

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -918,7 +918,7 @@ public abstract class kotlinx/coroutines/flow/internal/ChannelFlow : kotlinx/cor
 	public fun collect (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	protected abstract fun collectTo (Lkotlinx/coroutines/channels/ProducerScope;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	protected abstract fun create (Lkotlin/coroutines/CoroutineContext;I)Lkotlinx/coroutines/flow/internal/ChannelFlow;
-	public final fun produceImpl (Lkotlinx/coroutines/CoroutineScope;)Lkotlinx/coroutines/channels/ReceiveChannel;
+	public fun produceImpl (Lkotlinx/coroutines/CoroutineScope;)Lkotlinx/coroutines/channels/ReceiveChannel;
 	public fun toString ()Ljava/lang/String;
 	public final fun update (Lkotlin/coroutines/CoroutineContext;I)Lkotlinx/coroutines/flow/internal/ChannelFlow;
 	public static synthetic fun update$default (Lkotlinx/coroutines/flow/internal/ChannelFlow;Lkotlin/coroutines/CoroutineContext;IILjava/lang/Object;)Lkotlinx/coroutines/flow/internal/ChannelFlow;
@@ -929,6 +929,11 @@ public final class kotlinx/coroutines/flow/internal/SafeCollector : kotlinx/coro
 	public fun emit (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
+public final class kotlinx/coroutines/flow/internal/SendingCollector : kotlinx/coroutines/flow/internal/ConcurrentFlowCollector {
+	public fun <init> (Lkotlinx/coroutines/channels/SendChannel;)V
+	public fun emit (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public class kotlinx/coroutines/scheduling/ExperimentalCoroutineDispatcher : kotlinx/coroutines/ExecutorCoroutineDispatcher {
 	public synthetic fun <init> (II)V
 	public synthetic fun <init> (IIILkotlin/jvm/internal/DefaultConstructorMarker;)V

kotlinx-coroutines-core/common/src/flow/internal/ChannelFlow.kt

@@ -15,7 +15,11 @@ import kotlin.jvm.*
 internal fun <T> Flow<T>.asChannelFlow(): ChannelFlow<T> =
     this as? ChannelFlow ?: ChannelFlowOperatorImpl(this)
 
-// Operators that use channels extend this ChannelFlow and are always fused with each other
+/**
+ * Operators that use channels extend this ChannelFlow and are always fused with each other.
+ *
+ * @suppress **This an internal API and should not be used from general code.**
+ */
 @InternalCoroutinesApi
 public abstract class ChannelFlow<T>(
     // upstream context
@@ -63,7 +67,7 @@ public abstract class ChannelFlow<T>(
     fun broadcastImpl(scope: CoroutineScope, start: CoroutineStart): BroadcastChannel<T> =
         scope.broadcast(context, produceCapacity, start, block = collectToFun)
 
-    fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> =
+    open fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> =
         scope.flowProduce(context, produceCapacity, block = collectToFun)
 
     override suspend fun collect(collector: FlowCollector<T>) =

kotlinx-coroutines-core/common/src/flow/internal/Concurrent.kt

@@ -5,6 +5,7 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.atomicfu.*
+import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.channels.ArrayChannel
 import kotlinx.coroutines.flow.*
@@ -17,8 +18,13 @@ internal fun <T> FlowCollector<T>.asConcurrentFlowCollector(): ConcurrentFlowCol
 // Two basic implementations are here: SendingCollector and ConcurrentFlowCollector
 internal interface ConcurrentFlowCollector<T> : FlowCollector<T>
 
-// Concurrent collector because it sends to a channel
-internal class SendingCollector<T>(
+/**
+ * Collection that sends to channel. It is marked as [ConcurrentFlowCollector] because it can be used concurrently.
+ *
+ * @suppress **This an internal API and should not be used from general code.**
+ */
+@InternalCoroutinesApi
+public class SendingCollector<T>(
     private val channel: SendChannel<T>
 ) : ConcurrentFlowCollector<T> {
     override suspend fun emit(value: T) = channel.send(value)

kotlinx-coroutines-core/common/src/flow/internal/FlowCoroutine.kt

@@ -70,7 +70,6 @@ private class FlowCoroutine<T>(
     context: CoroutineContext,
     uCont: Continuation<T>
 ) : ScopeCoroutine<T>(context, uCont) {
-
     public override fun childCancelled(cause: Throwable): Boolean {
         if (cause is ChildCancelledException) return true
         return cancelImpl(cause)
@@ -81,7 +80,6 @@ private class FlowProduceCoroutine<T>(
     parentContext: CoroutineContext,
     channel: Channel<T>
 ) : ProducerCoroutine<T>(parentContext, channel) {
-
     public override fun childCancelled(cause: Throwable): Boolean {
         if (cause is ChildCancelledException) return true
         return cancelImpl(cause)

reactive/kotlinx-coroutines-reactive/src/Channel.kt

@@ -41,7 +41,7 @@ public suspend inline fun <T> Publisher<T>.collect(action: (T) -> Unit) =
     openSubscription().consumeEach(action)
 
 @Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
-private class SubscriptionChannel<T>(
+internal class SubscriptionChannel<T>(
     private val request: Int
 ) : LinkedListChannel<T>(), Subscriber<T> {
     init {

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

@@ -8,6 +8,7 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.flow.internal.*
+import kotlinx.coroutines.reactive.*
 import org.reactivestreams.*
 import kotlin.coroutines.*
 
@@ -34,84 +35,35 @@ public fun <T : Any> Publisher<T>.asFlow(): Flow<T> =
 )
 public fun <T : Any> Publisher<T>.asFlow(batchSize: Int): Flow<T> = asFlow().buffer(batchSize)
 
-
 private class PublisherAsFlow<T : Any>(
-    private val publisher: Publisher<T>, capacity: Int
-) : ChannelFlow<T>(
-    EmptyCoroutineContext,
-    capacity
-) {
-
-    override fun create(context: CoroutineContext, capacity: Int): ChannelFlow<T> {
-        return PublisherAsFlow(publisher, capacity)
-    }
-
-    /*
-     * It's possible to get rid of the second channel here, but it requires intrusive changes in ChannelFlow.
-     * Do it after API stabilization (including produceIn/broadcastIn).
-     */
-    override suspend fun collectTo(scope: ProducerScope<T>) = collect { scope.send(it) }
-
-    override suspend fun collect(collector: FlowCollector<T>) {
-        val channel = Channel<T>(capacity)
-        val subscriber = ReactiveSubscriber(channel, capacity)
-        publisher.subscribe(subscriber)
-        try {
-            var consumed = 0
-            for (i in channel) {
-                collector.emit(i)
-                if (++consumed == capacity) {
-                    consumed = 0
-                    subscriber.subscription.request(capacity.toLong())
-                }
-            }
-        } finally {
-            subscriber.subscription.cancel()
+    private val publisher: Publisher<T>,
+    capacity: Int
+) : ChannelFlow<T>(EmptyCoroutineContext, capacity) {
+    override fun create(context: CoroutineContext, capacity: Int): ChannelFlow<T> =
+        PublisherAsFlow(publisher, capacity)
+
+    override fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> {
+        val channel = publisher.openSubscription(capacity)
+        val handle = scope.coroutineContext[Job]?.invokeOnCompletion(onCancelling = true) { cause ->
+            channel.cancel(cause?.let {
+                it as? CancellationException ?: CancellationException("Job was cancelled", it)
+            })
         }
-    }
-
-    private suspend inline fun collect(emit: (element: T) -> Unit) {
-        val channel = Channel<T>(capacity)
-        val subscriber = ReactiveSubscriber(channel, capacity)
-        publisher.subscribe(subscriber)
-        try {
-            var consumed = 0
-            for (i in channel) {
-                emit(i)
-                if (++consumed == capacity) {
-                    consumed = 0
-                    subscriber.subscription.request(capacity.toLong())
-                }
+        if (handle != null && handle !== NonDisposableHandle) {
+            @Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+            (channel as SubscriptionChannel<*>).invokeOnClose {
+                handle.dispose()
             }
-        } finally {
-            subscriber.subscription.cancel()
         }
-    }
-}
-
-@Suppress("SubscriberImplementation")
-private class ReactiveSubscriber<T : Any>(
-    private val channel: Channel<T>,
-    private val requestSize: Int
-) : Subscriber<T> {
-
-    lateinit var subscription: Subscription
-
-    override fun onComplete() {
-        channel.close()
+        return channel
     }
 
-    override fun onSubscribe(s: Subscription) {
-        subscription = s
-        s.request(requestSize.toLong())
-    }
-
-    override fun onNext(t: T) {
-        // Controlled by requestSize
-        require(channel.offer(t)) { "Element $t was not added to channel because it was full, $channel" }
+    override suspend fun collect(collector: FlowCollector<T>) {
+        val channel = publisher.openSubscription(capacity)
+        channel.consumeEach { collector.emit(it) }
     }
 
-    override fun onError(t: Throwable?) {
-        channel.close(t)
-    }
+    // The second channel here is used only for broadcast
+    override suspend fun collectTo(scope: ProducerScope<T>) =
+        collect(SendingCollector(scope.channel))
 }

reactive/kotlinx-coroutines-reactive/test/flow/PublisherAsFlowTest.kt

@@ -11,7 +11,6 @@ import kotlinx.coroutines.reactive.*
 import kotlin.test.*
 
 class PublisherAsFlowTest : TestBase() {
-
     @Test
     fun testCancellation() = runTest {
         var onNext = 0
@@ -51,11 +50,11 @@ class PublisherAsFlowTest : TestBase() {
             send(3)
 
             expect(2)
-            send(5)
+            send(6)
 
             expect(4)
             send(7)
-            expect(6)
+            expect(5)
         }
 
         publisher.asFlow().collect {

reactive/kotlinx-coroutines-reactive/test/flow/RangePublisherBufferedTest.kt

@@ -10,8 +10,9 @@ import org.reactivestreams.*
 import org.reactivestreams.example.unicast.*
 import org.reactivestreams.tck.*
 
-class RangePublisherBufferedTest : PublisherVerification<Int>(TestEnvironment(50, 50)) {
-
+class RangePublisherBufferedTest :
+    PublisherVerification<Int>(TestEnvironment(50, 50))
+{
     override fun createPublisher(elements: Long): Publisher<Int> {
         return RangePublisher(1, elements.toInt()).asFlow().buffer(2).asPublisher()
     }