New flow builder: channelFlow (and its alias callbackFlow) and supplementary ProducerScope.await method

Rationale:

    * Can be used in different context without breaking context preservation
    * Can be used to build concurrent operators such as merge
    * Can be used to integrate with callbacks
    * Is less error-prone than flowViaChannel because requires explicit await() call

Partially fixes #1210

Author: qwwdfsad

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

@@ -719,6 +719,8 @@ public final class kotlinx/coroutines/channels/ConflatedBroadcastChannel : kotli
 }
 
 public final class kotlinx/coroutines/channels/ProduceKt {
+	public static final fun await (Lkotlinx/coroutines/channels/ProducerScope;Lkotlin/jvm/functions/Function0;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun await$default (Lkotlinx/coroutines/channels/ProducerScope;Lkotlin/jvm/functions/Function0;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static final fun produce (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/channels/ReceiveChannel;
 	public static final fun produce (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/channels/ReceiveChannel;
 	public static synthetic fun produce$default (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lkotlinx/coroutines/channels/ReceiveChannel;
@@ -797,6 +799,10 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun asFlow ([Ljava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun broadcastIn (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/CoroutineScope;ILkotlinx/coroutines/CoroutineStart;)Lkotlinx/coroutines/channels/BroadcastChannel;
 	public static synthetic fun broadcastIn$default (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/CoroutineScope;ILkotlinx/coroutines/CoroutineStart;ILjava/lang/Object;)Lkotlinx/coroutines/channels/BroadcastChannel;
+	public static final fun callbackFlow (ILkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static synthetic fun callbackFlow$default (ILkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun channelFlow (ILkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static synthetic fun channelFlow$default (ILkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun collect (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun combineLatest (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun count (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

kotlinx-coroutines-core/common/src/channels/Produce.kt

@@ -25,6 +25,32 @@ public interface ProducerScope<in E> : CoroutineScope, SendChannel<E> {
     val channel: SendChannel<E>
 }
 
+/**
+ * Suspends the current coroutine until the channel is either [closed][SendChannel.close] or [cancelled][ReceiveChannel.cancel]
+ * and invokes the given [block] before resuming the coroutine.
+ *
+ * Example of usage:
+ * ```
+ * val callbackEventsStream = produce {
+ *     val disposable = registerChannelInCallback(channel)
+ *     await { disposable.dispose() }
+ * }
+ * ```
+ */
+@ExperimentalCoroutinesApi
+public suspend fun <T> ProducerScope<T>.await(block: () -> Unit = {}) {
+    check(kotlin.coroutines.coroutineContext[Job] === this) { "await() can be invoke only from the producer context" }
+    suspendCancellableCoroutine<Unit> { cont ->
+        invokeOnClose {
+            try {
+                block()
+            } finally {
+                cont.resume(Unit)
+            }
+        }
+    }
+}
+
 /**
  * Launches new coroutine to produce a stream of values by sending them to a channel
  * and returns a reference to the coroutine as a [ReceiveChannel]. This resulting

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

@@ -188,38 +188,14 @@ public fun LongRange.asFlow(): Flow<Long> = flow {
 }
 
 /**
- * Creates an instance of the cold [Flow] with elements that are sent to a [SendChannel]
- * that is provided to the builder's [block] of code. It allows elements to be
- * produced by the code that is running in a different context, e.g. from a callback-based API.
- *
- * The resulting flow is _cold_, which means that [block] is called on each call of a terminal operator
- * on the resulting flow. The [block] is not suspending deliberately, if you need suspending scope, [flow] builder
- * should be used instead.
- *
- * To control backpressure, [bufferSize] is used and matches directly the `capacity` parameter of [Channel] factory.
- * The provided channel can later be used by any external service to communicate with flow and its buffer determines
- * backpressure buffer size or its behaviour (e.g. in case when [Channel.CONFLATED] was used).
- *
- * Example of usage:
- * ```
- * fun flowFrom(api: CallbackBasedApi): Flow<T> = flowViaChannel { channel ->
- *     val callback = object : Callback { // implementation of some callback interface
- *         override fun onNextValue(value: T) {
- *             channel.offer(value) // Note: offer drops value when buffer is full
- *         }
- *         override fun onApiError(cause: Throwable) {
- *             channel.cancel("API Error", CancellationException(cause))
- *         }
- *         override fun onCompleted() = channel.close()
- *     }
- *     api.register(callback)
- *     channel.invokeOnClose {
- *         api.unregister(callback)
- *     }
- * }
- * ```
+ * @suppress
  */
 @FlowPreview
+@Deprecated(
+    message = "Use channelFlow instead",
+    level = DeprecationLevel.WARNING,
+    replaceWith = ReplaceWith("channelFlow(bufferSize, block)")
+)
 public fun <T> flowViaChannel(
     bufferSize: Int = 16,
     @BuilderInference block: CoroutineScope.(channel: SendChannel<T>) -> Unit
@@ -237,3 +213,88 @@ public fun <T> flowViaChannel(
         }
     }
 }
+
+/**
+ * Creates an instance of the cold [Flow] with elements that are sent to a [SendChannel]
+ * that is provided to the builder's [block] of code via [ProducerScope]. It allows elements to be
+ * produced by the code that is running in a different context or running concurrently.
+ * The resulting flow is _cold_, which means that [block] is called on each call of a terminal operator
+ * on the resulting flow.
+ *
+ * This builder ensures thread-safety and context preservation, thus the provided [ProducerScope] can be used concurrently from different contexts.
+ * The resulting flow will complete as soon as [ProducerScope], to artificially prolong it [await] can be used.
+ * For more detailed example please refer to [callbackFlow] documentation.
+ *
+ * To control backpressure, [bufferSize] is used and matches directly the `capacity` parameter of [Channel] factory.
+ * The provided channel can later be used by any external service to communicate with the flow and its buffer determines
+ * backpressure buffer size or its behaviour (e.g. in the case when [Channel.CONFLATED] was used).
+ *
+ * Examples of usage:
+ * ```
+ * fun <T> Flow<T>.merge(other: Flow<T>): Flow<T> = channelFlow {
+ *     launch {
+ *         collect { value -> send(value) }
+ *     }
+ *     other.collect { value -> send(value) }
+ * }
+ *
+ * fun <T> contextualFlow(): Flow<T> = channelFlow {
+ *     launch(Dispatchers.IO) {
+ *         send(computeIoValue())
+ *     }
+ *
+ *     launch(Dispatchers.Default) {
+ *         send(computeCpuValue())
+ *     }
+ * }
+ * ```
+ */
+@FlowPreview
+public fun <T> channelFlow(bufferSize: Int = 16, @BuilderInference block: suspend ProducerScope<T>.() -> Unit): Flow<T> =
+    flow {
+        coroutineScope {
+            val channel = produce(capacity = bufferSize, block = block)
+            channel.consumeEach { value ->
+                emit(value)
+            }
+        }
+    }
+
+/**
+ * Creates an instance of the cold [Flow] with elements that are sent to a [SendChannel]
+ * that is provided to the builder's [block] of code via [ProducerScope]. It allows elements to be
+ * produced by the code that is running in a different context or running concurrently.
+ *
+ * The resulting flow is _cold_, which means that [block] is called on each call of a terminal operator
+ * on the resulting flow.
+ *
+ * This builder ensures thread-safety and context preservation, thus the provided [ProducerScope] can be used from any context,
+ * e.g. from the callback-based API. The flow completes as soon as its scope completes, thus if you are using channel from the
+ * callback-based API, to artificially prolong scope lifetime and avoid memory-leaks related to unregistered resources,
+ * [await] extension should be used. [await] argument will be invoked when either flow consumer cancels flow collection
+ * or when callback-based API invokes [SendChannel.close] manually.
+ *
+ * To control backpressure, [bufferSize] is used and matches directly the `capacity` parameter of [Channel] factory.
+ * The provided channel can later be used by any external service to communicate with the flow and its buffer determines
+ * backpressure buffer size or its behaviour (e.g. in the case when [Channel.CONFLATED] was used).
+ *
+ * Example of usage:
+ * ```
+ * fun flowFrom(api: CallbackBasedApi): Flow<T> = callbackFlow {
+ *     val callback = object : Callback { // implementation of some callback interface
+ *         override fun onNextValue(value: T) {
+ *             offer(value) // Note: offer drops value when buffer is full
+ *         }
+ *         override fun onApiError(cause: Throwable) {
+ *             cancel("API Error", CancellationException(cause))
+ *         }
+ *         override fun onCompleted() = channel.close()
+ *     }
+ *     api.register(callback)
+ *     // Suspend until either onCompleted or external cancellation are invoked
+ *     await { api.unregister(callback) }
+ * }
+ * ```
+ */
+public inline fun <T> callbackFlow(bufferSize: Int = 16, @BuilderInference crossinline block: suspend ProducerScope<T>.() -> Unit): Flow<T> =
+    channelFlow(bufferSize) { block() }

kotlinx-coroutines-core/common/test/channels/ProduceTest.kt

@@ -94,6 +94,60 @@ class ProduceTest : TestBase() {
         cancelOnCompletion(coroutineContext)
     }
 
+    @Test
+    fun testAwaitConsumerCancellation() = runTest {
+        val parent = Job()
+        val channel = produce<Int>(parent) {
+            expect(2)
+            await { expect(4) }
+        }
+        expect(1)
+        yield()
+        expect(3)
+        channel.cancel()
+        parent.complete()
+        parent.join()
+        finish(5)
+    }
+
+    @Test
+    fun testAwaitProducerCancellation() = runTest {
+        val parent = Job()
+        produce<Int>(parent) {
+            expect(2)
+            launch {
+                expect(3)
+                this@produce.cancel()
+            }
+            await { expect(4) }
+        }
+        expect(1)
+        parent.complete()
+        parent.join()
+        finish(5)
+    }
+
+    @Test
+    fun testAwaitParentCancellation() = runTest {
+        val parent = Job()
+        produce<Int>(parent) {
+            expect(2)
+            await { expect(4) }
+        }
+        expect(1)
+        yield()
+        expect(3)
+        parent.cancelAndJoin()
+        finish(5)
+    }
+
+    @Test
+    fun testAwaitIllegalState() = runTest {
+        val channel = produce<Int> {  }
+        @Suppress("RemoveExplicitTypeArguments") // KT-31525
+        assertFailsWith<IllegalStateException> { (channel as ProducerScope<*>).await<Nothing>() }
+    }
+
     private suspend fun cancelOnCompletion(coroutineContext: CoroutineContext) = CoroutineScope(coroutineContext).apply {
         val source = Channel<Int>()
         expect(1)

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

@@ -0,0 +1,150 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import kotlin.test.*
+
+class ChannelBuildersFlowTest : TestBase() {
+    @Test
+    fun testBroadcastChannelAsFlow() = runTest {
+        val channel = broadcast {
+           repeat(10) {
+               send(it + 1)
+           }
+        }
+
+        val sum = channel.asFlow().sum()
+        assertEquals(55, sum)
+    }
+
+    @Test
+    fun testExceptionInBroadcast() = runTest {
+        expect(1)
+        val channel = broadcast(NonCancellable) { // otherwise failure will cancel scope as well
+            repeat(10) {
+                send(it + 1)
+            }
+            throw TestException()
+        }
+        assertEquals(15, channel.asFlow().take(5).sum())
+
+        // Workaround for JS bug
+        try {
+            channel.asFlow().collect { /* Do nothing */ }
+            expectUnreached()
+        } catch (e: TestException) {
+            finish(2)
+        }
+    }
+
+    @Test
+    fun testBroadcastChannelAsFlowLimits() = runTest {
+        val channel = BroadcastChannel<Int>(1)
+        val flow = channel.asFlow().map { it * it }.drop(1).take(2)
+
+        var expected = 0
+        launch {
+            assertTrue(channel.offer(1)) // Handed to the coroutine
+            assertTrue(channel.offer(2)) // Buffered
+            assertFalse(channel.offer(3)) // Failed to offer
+            channel.send(3)
+            yield()
+            assertEquals(1, expected)
+            assertTrue(channel.offer(4)) // Handed to the coroutine
+            assertTrue(channel.offer(5)) // Buffered
+            assertFalse(channel.offer(6))  // Failed to offer
+            channel.send(6)
+            assertEquals(2, expected)
+        }
+
+        val sum = flow.sum()
+        assertEquals(13, sum)
+        ++expected
+        val sum2 = flow.sum()
+        assertEquals(61, sum2)
+        ++expected
+    }
+
+    @Test
+    fun flowAsBroadcast() = runTest {
+        val flow = flow {
+            repeat(10) {
+                emit(it)
+            }
+        }
+
+        val channel = flow.broadcastIn(this)
+        assertEquals((0..9).toList(), channel.openSubscription().toList())
+    }
+
+    @Test
+    fun flowAsBroadcastMultipleSubscription() = runTest {
+        val flow = flow {
+            repeat(10) {
+                emit(it)
+            }
+        }
+
+        val broadcast = flow.broadcastIn(this)
+        val channel = broadcast.openSubscription()
+        val channel2 = broadcast.openSubscription()
+
+        assertEquals(0, channel.receive())
+        assertEquals(0, channel2.receive())
+        yield()
+        assertEquals(1, channel.receive())
+        assertEquals(1, channel2.receive())
+
+        channel.cancel()
+        channel2.cancel()
+        yield()
+        ensureActive()
+    }
+
+    @Test
+    fun flowAsBroadcastException() = runTest {
+        val flow = flow {
+            repeat(10) {
+                emit(it)
+            }
+
+            throw TestException()
+        }
+
+        val channel = flow.broadcastIn(this + NonCancellable)
+        assertFailsWith<TestException> { channel.openSubscription().toList() }
+        assertTrue(channel.isClosedForSend) // Failure in the flow fails the channel
+    }
+
+    // Semantics of these tests puzzle me, we should figure out the way to prohibit such chains
+    @Test
+    fun testFlowAsBroadcastAsFlow() = runTest {
+        val flow = flow {
+            emit(1)
+            emit(2)
+            emit(3)
+        }.broadcastIn(this).asFlow()
+
+        assertEquals(6, flow.sum())
+        assertEquals(0, flow.sum()) // Well suddenly flow is no longer idempotent and cold
+    }
+
+    @Test
+    fun testBroadcastAsFlowAsBroadcast() = runTest {
+        val channel = broadcast {
+            send(1)
+        }.asFlow().broadcastIn(this)
+
+        channel.openSubscription().consumeEach {
+            assertEquals(1, it)
+        }
+
+        channel.openSubscription().consumeEach {
+            fail()
+        }
+    }
+}