Incorporate flow scope into flow operators

Author: qwwdfsad

kotlinx-coroutines-core/common/src/JobSupport.kt

@@ -620,16 +620,15 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
     }
 
     /**
-     * Returns `true` if job should cancel itself during on child [CancellationException].
+     * Returns `true` if job should cancel itself on child [CancellationException].
      */
     public open fun cancelOnChildCancellation(cause: CancellationException) = false
 
     /**
      * Child was cancelled with a cause.
-     * At this point parent decides whether it cancels itself (e.g. on a critical failure) and
+     * In this method parent decides whether it cancels itself (e.g. on a critical failure) and
      * whether it handles the exception of the child.
-     *
-     * It is overridden in supervisor implementations to completely ignore any child cancellation
+     * It is overridden in supervisor implementations to completely ignore any child cancellation.
      */
     public open fun childCancelled(cause: Throwable): Boolean {
         if (cause is CancellationException && !cancelOnChildCancellation(cause)) return true
@@ -949,11 +948,8 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
 
     /**
      * Returns `true` for scoped coroutines.
-     * Scoped coroutine is a coroutine that is executed sequentially within the enclosing scope
-     * without any concurrency.
-     * Scoped coroutines always handle any exception happened within -- they just rethrow it
-     * to the enclosing scope.
-     *
+     * Scoped coroutine is a coroutine that is executed sequentially within the enclosing scope without any concurrency.
+     * Scoped coroutines always handle any exception happened within -- they just rethrow it to the enclosing scope.
      * Examples of scoped coroutines are `coroutineScope`, `withTimeout` and `runBlocking`.
      */
     protected open val isScopedCoroutine: Boolean get() = false

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

@@ -118,7 +118,6 @@ public fun <T> Flow<T>.onErrorResume(fallback: Flow<T>): Flow<T> = error("Should
 @Deprecated(message = "withContext in flow body is deprecated, use flowOn instead", level = DeprecationLevel.ERROR)
 public fun <T, R> FlowCollector<T>.withContext(context: CoroutineContext, block: suspend () -> R): Unit = error("Should not be called")
 
-
 /**
  * `subscribe` is Rx-specific API that has no direct match in flows.
  * One can use `launch` instead, for example the following:

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

@@ -5,33 +5,56 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
 /**
- * Creates a [CoroutineScope] and calls the specified suspend block with this scope.
- * This builder is similar to [coroutineScope] with the only exception that it *ties* lifecycle of children
- * and itself regarding the cancellation, thus being cancelled when one of the children becomes cancelled.
+ * Creates a flow that also provides a [CoroutineScope] for each collector.
  *
+ * Provided [CoroutineScope] is similar to [coroutineScope] with the only exception that it *ties* lifecycle of children
+ * and itself regarding the cancellation, thus being cancelled when one of the children becomes cancelled.
  * For example:
  * ```
- * flowScope {
+ * scopedFlow {
  *     launch {
  *         throw CancellationException()
  *     }
  * } // <- CE will be rethrown here
  * ```
+ *
+ * Basically, it is shorthand for:
+ * ```
+ * flow {
+ *     coroutineScope {
+ *         ...
+ *     }
+ * }
+ * ```
+ * with additional constraint on cancellation.
+ * To cancel child without cancelling itself, `cancel(ChildCancelledException())` should be used.
  */
-internal suspend fun <R> flowScope(block: suspend CoroutineScope.() -> R): R =
-    suspendCoroutineUninterceptedOrReturn { uCont ->
-        val coroutine = FlowScope(uCont.context, uCont)
-        coroutine.startUndispatchedOrReturn(coroutine, block)
+internal fun <R> scopedFlow(@BuilderInference block: suspend FlowScope<Unit, R>.() -> Unit): Flow<R> =
+    object : Flow<R> {
+        override suspend fun collect(collector: FlowCollector<R>) {
+            suspendCoroutineUninterceptedOrReturn<Unit> { uCont ->
+                val coroutine = FlowScope(collector, uCont.context, uCont)
+                coroutine.startUndispatchedOrReturn(coroutine, block)
+            }
+        }
     }
 
-private class FlowScope<T>(context: CoroutineContext, uCont: Continuation<T>) : ScopeCoroutine<T>(context, uCont) {
+internal class FlowScope<T, R>(
+    private val collector: FlowCollector<R>,
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(context, uCont), FlowCollector<R> {
+
     public override fun cancelOnChildCancellation(cause: CancellationException) = cause !is ChildCancelledException
+
+    override suspend fun emit(value: R) = collector.emit(value)
 }
 
 internal class ChildCancelledException : CancellationException(null)

kotlinx-coroutines-core/common/src/flow/operators/Delay.kt

@@ -60,34 +60,33 @@ public fun <T> Flow<T>.delayEach(timeMillis: Long): Flow<T> = flow {
  */
 public fun <T> Flow<T>.debounce(timeoutMillis: Long): Flow<T> {
     require(timeoutMillis > 0) { "Debounce timeout should be positive" }
-    return flow {
-        coroutineScope {
-            val values = Channel<Any?>(Channel.CONFLATED) // Actually Any, KT-30796
-            // Channel is not closed deliberately as there is no close with value
-            val collector = async {
-                collect { value -> values.send(value ?: NULL) }
-            }
+    return scopedFlow {
+        val values = Channel<Any?>(Channel.CONFLATED) // Actually Any, KT-30796
+        // Channel is not closed deliberately as there is no close with value
+        val collector = async {
+            collect { value -> values.send(value ?: NULL) }
+        }
 
-            var isDone = false
-            var lastValue: Any? = null
-            while (!isDone) {
-                select<Unit> {
-                    values.onReceive {
-                        lastValue = it
-                    }
+        var isDone = false
+        var lastValue: Any? = null
+        while (!isDone) {
+            select<Unit> {
+                values.onReceive {
+                    lastValue = it
+                }
 
-                    lastValue?.let { value -> // set timeout when lastValue != null
-                        onTimeout(timeoutMillis) {
-                            lastValue = null // Consume the value
-                            emit(NULL.unbox(value))
-                        }
+                lastValue?.let { value ->
+                    // set timeout when lastValue != null
+                    onTimeout(timeoutMillis) {
+                        lastValue = null // Consume the value
+                        emit(NULL.unbox(value))
                     }
+                }
 
-                    // Close with value 'idiom'
-                    collector.onAwait {
-                        if (lastValue != null) emit(NULL.unbox(lastValue))
-                        isDone = true
-                    }
+                // Close with value 'idiom'
+                collector.onAwait {
+                    if (lastValue != null) emit(NULL.unbox(lastValue))
+                    isDone = true
                 }
             }
         }
@@ -112,32 +111,31 @@ public fun <T> Flow<T>.debounce(timeoutMillis: Long): Flow<T> {
  */
 public fun <T> Flow<T>.sample(periodMillis: Long): Flow<T> {
     require(periodMillis > 0) { "Sample period should be positive" }
-    return flow {
-        coroutineScope {
-            val values = produce<Any?>(capacity = Channel.CONFLATED) {  // Actually Any, KT-30796
-                collect { value -> send(value ?: NULL) }
-            }
+    return scopedFlow {
+        val values = produce<Any?>(capacity = Channel.CONFLATED) {
+            // Actually Any, KT-30796
+            collect { value -> send(value ?: NULL) }
+        }
 
-            var isDone = false
-            var lastValue: Any? = null
-            val ticker = fixedPeriodTicker(periodMillis)
-            while (!isDone) {
-                select<Unit> {
-                    values.onReceiveOrNull {
-                        if (it == null) {
-                            ticker.cancel()
-                            isDone = true
-                        } else {
-                            lastValue = it
-                        }
+        var isDone = false
+        var lastValue: Any? = null
+        val ticker = fixedPeriodTicker(periodMillis)
+        while (!isDone) {
+            select<Unit> {
+                values.onReceiveOrNull {
+                    if (it == null) {
+                        ticker.cancel(ChildCancelledException())
+                        isDone = true
+                    } else {
+                        lastValue = it
                     }
+                }
 
-                    // todo: shall be start sampling only when an element arrives or sample aways as here?
-                    ticker.onReceive {
-                        val value = lastValue ?: return@onReceive
-                        lastValue = null // Consume the value
-                        emit(NULL.unbox(value))
-                    }
+                // todo: shall be start sampling only when an element arrives or sample aways as here?
+                ticker.onReceive {
+                    val value = lastValue ?: return@onReceive
+                    lastValue = null // Consume the value
+                    emit(NULL.unbox(value))
                 }
             }
         }

kotlinx-coroutines-core/common/src/flow/operators/Merge.kt

@@ -129,17 +129,16 @@ public fun <T> Flow<Flow<T>>.flattenMerge(concurrency: Int = DEFAULT_CONCURRENCY
  * produces `aa bb b_last`
  */
 @FlowPreview
-public fun <T, R> Flow<T>.switchMap(transform: suspend (value: T) -> Flow<R>): Flow<R> = flow {
-    coroutineScope {
-        var previousFlow: Job? = null
-        collect { value ->
-            // Linearize calls to emit as alternative to the channel. Bonus points for never-overlapping channels.
-            previousFlow?.cancelAndJoin()
-            // Undispatched to have better user experience in case of synchronous flows
-            previousFlow = launch(start = CoroutineStart.UNDISPATCHED) {
-                transform(value).collect { innerValue ->
-                    emit(innerValue)
-                }
+public fun <T, R> Flow<T>.switchMap(transform: suspend (value: T) -> Flow<R>): Flow<R> = scopedFlow {
+    var previousFlow: Job? = null
+    collect { value ->
+        // Linearize calls to emit as alternative to the channel. Bonus points for never-overlapping channels.
+        previousFlow?.cancel(ChildCancelledException())
+        previousFlow?.join()
+        // Undispatched to have better user experience in case of synchronous flows
+        previousFlow = launch(start = CoroutineStart.UNDISPATCHED) {
+            transform(value).collect { innerValue ->
+                emit(innerValue)
             }
         }
     }

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

@@ -81,4 +81,35 @@ class ChannelFlowTest : TestBase() {
         assertFailsWith<TestException>(flow)
         finish(4)
     }
+
+    @Test
+    fun testMergeOneCoroutineWithCancellation() = runTest {
+        val flow = flowOf(1, 2, 3)
+        val f = flow.mergeOneCoroutine(flow).take(2)
+        assertEquals(listOf(1, 1), f.toList())
+    }
+
+    @Test
+    fun testMergeTwoCoroutinesWithCancellation() = runTest {
+        val flow = flowOf(1, 2, 3)
+        val f = flow.mergeTwoCoroutines(flow).take(2)
+        assertEquals(listOf(1, 1), f.toList())
+    }
+
+    private fun Flow<Int>.mergeTwoCoroutines(other: Flow<Int>): Flow<Int> = channelFlow {
+        launch {
+            collect { send(it); yield() }
+        }
+        launch {
+            other.collect { send(it) }
+        }
+    }
+
+    private fun Flow<Int>.mergeOneCoroutine(other: Flow<Int>): Flow<Int> = channelFlow {
+        launch {
+            collect { send(it); yield() }
+        }
+
+        other.collect { send(it); yield() }
+    }
 }

kotlinx-coroutines-core/common/test/flow/internal/FlowScopeTest.kt

@@ -5,10 +5,15 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 import kotlin.test.*
 
 class FlowScopeTest : TestBase() {
 
+    private suspend fun flowScope(block: suspend CoroutineScope.() -> Unit): Unit {
+        scopedFlow<Unit>(block).singleOrNull()
+    }
+
     @Test
     fun testCancellation() = runTest {
         assertFailsWith<CancellationException> {

kotlinx-coroutines-core/common/test/flow/operators/CombineLatestTest.kt

@@ -197,6 +197,46 @@ abstract class CombineLatestTestBase : TestBase() {
         assertFailsWith<TestException>(flow)
         finish(2)
     }
+
+    @Test
+    fun testCancellationExceptionUpstream() = runTest {
+        val f1 = flow {
+            expect(1)
+            emit(1)
+            throw CancellationException("")
+        }
+        val f2 = flow {
+            emit(1)
+            hang { expect(3) }
+        }
+
+        val flow = f1.combineLatest(f2, { _, _ -> 1 }).onEach { expect(2) }
+        assertFailsWith<CancellationException>(flow)
+        finish(4)
+    }
+
+    @Test
+    fun testCancellationExceptionDownstream() = runTest {
+        val f1 = flow {
+            emit(1)
+            expect(2)
+            hang { expect(5) }
+        }
+        val f2 = flow {
+            emit(1)
+            expect(3)
+            hang { expect(6) }
+        }
+
+        val flow = f1.combineLatest(f2, { _, _ -> 1 }).onEach {
+            expect(1)
+            yield()
+            expect(4)
+            throw CancellationException("")
+        }
+        assertFailsWith<CancellationException>(flow)
+        finish(7)
+    }
 }
 
 class CombineLatestTest : CombineLatestTestBase() {

kotlinx-coroutines-core/common/test/flow/operators/DebounceTest.kt

@@ -95,20 +95,25 @@ class DebounceTest : TestBase() {
     }
 
     @Test
-    fun testUpstreamError() = runTest {
+    fun testUpstreamError()= testUpstreamError(TimeoutCancellationException(""))
+
+    @Test
+    fun testUpstreamErrorCancellation() = testUpstreamError(TimeoutCancellationException(""))
+
+    private inline fun <reified T: Throwable> testUpstreamError(cause: T) = runTest {
         val latch = Channel<Unit>()
         val flow = flow {
             expect(1)
             emit(1)
             expect(2)
             latch.receive()
-            throw TestException()
+            throw cause
         }.debounce(1).map {
             latch.send(Unit)
             hang { expect(3) }
         }
 
-        assertFailsWith<TestException>(flow)
+        assertFailsWith<T>(flow)
         finish(4)
     }