Introduce ReceiveChannel.consumeAsFlow and FlowCollector.emitAll(chan)

* This is a consuming conversion -- the resulting flow can be collected
  just once and the channel is closed after the first collect.
* The implementation is made efficient via emitAll extension.
* Experimental FlowCollector.emitAll extension is introduced.
* It is based on the (internal) Channel.receiveOrClose
  and ensures that the reference to the last emitted value is
  not retained (does not leak).

Fixes #1340
Fixes #1333

Author: elizarov

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

@@ -856,6 +856,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun concatWith (Lkotlinx/coroutines/flow/Flow;Ljava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun concatWith (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun conflate (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun consumeAsFlow (Lkotlinx/coroutines/channels/ReceiveChannel;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun count (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun count (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun debounce (Lkotlinx/coroutines/flow/Flow;J)Lkotlinx/coroutines/flow/Flow;
@@ -866,6 +867,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun distinctUntilChangedBy (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun drop (Lkotlinx/coroutines/flow/Flow;I)Lkotlinx/coroutines/flow/Flow;
 	public static final fun dropWhile (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun emitAll (Lkotlinx/coroutines/flow/FlowCollector;Lkotlinx/coroutines/channels/ReceiveChannel;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun emitAll (Lkotlinx/coroutines/flow/FlowCollector;Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun emptyFlow ()Lkotlinx/coroutines/flow/Flow;
 	public static final fun filter (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;

kotlinx-coroutines-core/common/src/flow/Channels.kt

@@ -7,14 +7,81 @@
 
 package kotlinx.coroutines.flow
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
-import kotlinx.coroutines.channels.Channel.Factory.CONFLATED
-import kotlinx.coroutines.channels.Channel.Factory.BUFFERED
-import kotlinx.coroutines.channels.Channel.Factory.OPTIONAL_CHANNEL
 import kotlinx.coroutines.flow.internal.*
-import kotlin.coroutines.*
 import kotlin.jvm.*
+import kotlinx.coroutines.flow.unsafeFlow as flow
+
+/**
+ * Emits all elements from the given [channel] to this flow collector and [cancels][cancel] (consumes)
+ * the channel afterwards. If you need to iterate over the channel without consuming it,
+ * a regular `for` loop should be used instead.
+ *
+ * This function provides a more efficient shorthand for `channel.consumeEach { value -> emit(value) }`.
+ * See [consumeEach][ReceiveChannel.consumeEach].
+ */
+@ExperimentalCoroutinesApi
+public suspend fun <T> FlowCollector<T>.emitAll(channel: ReceiveChannel<T>) {
+    // Manually inlined "consumeEach" implementation that does not use iterator but works via "receiveOrClosed".
+    // It has smaller and more efficient spilled state which also allows to implement a manual kludge to
+    // fix retention of the last emitted value.
+    // See https://youtrack.jetbrains.com/issue/KT-16222
+    // See https://github.com/Kotlin/kotlinx.coroutines/issues/1333
+    var cause: Throwable? = null
+    try {
+        while (true) {
+            // :KLUDGE: This "run" call is resolved to an extension function "run" and forces the size of
+            // spilled state to increase by an additional slot, so there are 4 object local variables spilled here
+            // which makes the size of spill state equal to the 4 slots that are spilled around subsequent "emit"
+            // call, ensuring that the previously emitted value is not retained in the state while receiving
+            // the next one.
+            //     L$0 <- this
+            //     L$1 <- channel
+            //     L$2 <- cause
+            //     L$3 <- this$run (actually equal to this)
+            val result = run { channel.receiveOrClosed() }
+            if (result.isClosed) {
+                result.closeCause?.let { throw it }
+                break // returns normally when result.closeCause == null
+            }
+            // result is spilled here to the coroutine state and retained after the call, even though
+            // it is not actually needed in the next loop iteration.
+            //     L$0 <- this
+            //     L$1 <- channel
+            //     L$2 <- cause
+            //     L$3 <- result
+            emit(result.value)
+        }
+    } catch (e: Throwable) {
+        cause = e
+        throw e
+    } finally {
+        channel.cancelConsumed(cause)
+    }
+}
+
+/**
+ * Represents the given receive channel as a hot flow and [consumes][ReceiveChannel.consume] the channel
+ * on the first collection from this flow. The resulting flow can be collected just once and throws
+ * [IllegalStateException] when trying to collect it more than once.
+ *
+ * ### Cancellation semantics
+ * 1) Flow consumer is cancelled when the original channel is cancelled.
+ * 2) Flow consumer completes normally when the original channel completes (~is closed) normally.
+ * 3) If the flow consumer fails with an exception, channel is cancelled.
+ *
+ */
+@FlowPreview
+public fun <T> ReceiveChannel<T>.consumeAsFlow(): Flow<T> = object : Flow<T> {
+    val collected = atomic(false)
+
+    override suspend fun collect(collector: FlowCollector<T>) {
+        check(!collected.getAndSet(true)) { "ReceiveChannel.consumeAsFlow can be collected just once" }
+        collector.emitAll(this@consumeAsFlow)
+    }
+}
 
 /**
  * Represents the given broadcast channel as a hot flow.
@@ -27,10 +94,7 @@ import kotlin.jvm.*
  */
 @FlowPreview
 public fun <T> BroadcastChannel<T>.asFlow(): Flow<T> = flow {
-    val subscription = openSubscription()
-    subscription.consumeEach { value ->
-        emit(value)
-    }
+    emitAll(openSubscription())
 }
 
 /**

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

@@ -72,8 +72,7 @@ public abstract class ChannelFlow<T>(
 
     override suspend fun collect(collector: FlowCollector<T>) =
         coroutineScope {
-            val channel = produceImpl(this)
-            channel.consumeEach { collector.emit(it) }
+            collector.emitAll(produceImpl(this))
         }
 
     // debug toString

kotlinx-coroutines-core/common/test/flow/channels/ChannelBuildersFlowTest.kt

@@ -9,6 +9,46 @@ import kotlinx.coroutines.channels.*
 import kotlin.test.*
 
 class ChannelBuildersFlowTest : TestBase() {
+    @Test
+    fun testChannelConsumeAsFlow() = runTest {
+        val channel = produce {
+           repeat(10) {
+               send(it + 1)
+           }
+        }
+        val flow = channel.consumeAsFlow()
+        assertEquals(55, flow.sum())
+        assertFailsWith<IllegalStateException> { flow.collect() }
+    }
+
+    @Test
+    fun testConsumeAsFlowCancellation() = runTest {
+        val channel = produce(NonCancellable) { // otherwise failure will cancel scope as well
+            repeat(10) {
+                send(it + 1)
+            }
+            throw TestException()
+        }
+        val flow = channel.consumeAsFlow()
+        assertEquals(15, flow.take(5).sum())
+        // the channel should have been canceled, even though took only 5 elements
+        assertTrue(channel.isClosedForReceive)
+        assertFailsWith<IllegalStateException> { flow.collect() }
+    }
+
+    @Test
+    fun testConsumeAsFlowException() = runTest {
+        val channel = produce(NonCancellable) { // otherwise failure will cancel scope as well
+            repeat(10) {
+                send(it + 1)
+            }
+            throw TestException()
+        }
+        val flow = channel.consumeAsFlow()
+        assertFailsWith<TestException> { flow.sum() }
+        assertFailsWith<IllegalStateException> { flow.collect() }
+    }
+
     @Test
     fun testBroadcastChannelAsFlow() = runTest {
         val channel = broadcast {