Check for cancellation in concurrent flow merge on each element

  * Implementation detail (launch on each value) is leaking into upstream behaviour
  * The overhead is negligible compared to launching a new coroutines and sending to channel, but it provides a much approachable mental model when no suspension in the upstream flow happens (note: upstream never sends elements to the channel)

Fixes #1392

4 files changed, 51 insertions, 1 deletions

diff --git a/kotlinx-coroutines-core/common/src/flow/operators/Merge.kt b/kotlinx-coroutines-core/common/src/flow/operators/Merge.kt
index e593d035..b1fe91ab 100644
--- a/kotlinx-coroutines-core/common/src/flow/operators/Merge.kt
+++ b/kotlinx-coroutines-core/common/src/flow/operators/Merge.kt
@@ -151,8 +151,14 @@ private class ChannelFlowMerge<T>(
     // The actual merge implementation with concurrency limit
     private suspend fun mergeImpl(scope: CoroutineScope, collector: ConcurrentFlowCollector<T>) {
         val semaphore = Semaphore(concurrency)
-        @Suppress("UNCHECKED_CAST")
+        val job: Job? = coroutineContext[Job]
         flow.collect { inner ->
+            /*
+             * We launch a coroutine on each emitted element and the only potential
+             * suspension point in this collector is `semaphore.acquire` that rarely suspends,
+             * so we manually check for cancellation to propagate it to the upstream in time.
+             */
+            job?.ensureActive()
             semaphore.acquire() // Acquire concurrency permit
             scope.launch {
                 try {
diff --git a/kotlinx-coroutines-core/common/test/flow/operators/BufferTest.kt b/kotlinx-coroutines-core/common/test/flow/operators/BufferTest.kt
index 65fef02c..0b1b208f 100644
--- a/kotlinx-coroutines-core/common/test/flow/operators/BufferTest.kt
+++ b/kotlinx-coroutines-core/common/test/flow/operators/BufferTest.kt
@@ -183,5 +183,19 @@ class BufferTest : TestBase() {
             }
         finish(n + 4)
     }
+
+    @Test
+    fun testCancellation() = runTest {
+        val result = flow {
+            emit(1)
+            emit(2)
+            emit(3)
+            expectUnreached()
+            emit(4)
+        }.buffer(0)
+            .take(2)
+            .toList()
+        assertEquals(listOf(1, 2), result)
+    }
 }
 
diff --git a/kotlinx-coroutines-core/common/test/flow/operators/FlatMapMergeTest.kt b/kotlinx-coroutines-core/common/test/flow/operators/FlatMapMergeTest.kt
index 6069ae6d..511a003a 100644
--- a/kotlinx-coroutines-core/common/test/flow/operators/FlatMapMergeTest.kt
+++ b/kotlinx-coroutines-core/common/test/flow/operators/FlatMapMergeTest.kt
@@ -73,4 +73,19 @@ class FlatMapMergeTest : FlatMapMergeBaseTest() {
         assertFailsWith<CancellationException>(flow)
         finish(5)
     }
+
+    @Test
+    fun testCancellation() = runTest {
+        val result = flow {
+            emit(1)
+            emit(2)
+            emit(3)
+            emit(4)
+            expectUnreached() // Cancelled by take
+            emit(5)
+        }.flatMapMerge(2) { v -> flow { emit(v) } }
+            .take(2)
+            .toList()
+        assertEquals(listOf(1, 2), result)
+    }
 }
diff --git a/kotlinx-coroutines-core/common/test/flow/operators/FlowOnTest.kt b/kotlinx-coroutines-core/common/test/flow/operators/FlowOnTest.kt
index 5d5fa6c6..34c0476e 100644
--- a/kotlinx-coroutines-core/common/test/flow/operators/FlowOnTest.kt
+++ b/kotlinx-coroutines-core/common/test/flow/operators/FlowOnTest.kt
@@ -262,6 +262,21 @@ class FlowOnTest : TestBase() {
     }
 
     @Test
+    fun testCancellation() = runTest {
+        val result = flow {
+            emit(1)
+            emit(2)
+            emit(3)
+            expectUnreached()
+            emit(4)
+        }.flowOn(wrapperDispatcher())
+            .buffer(0)
+            .take(2)
+            .toList()
+        assertEquals(listOf(1, 2), result)
+    }
+
+    @Test
     fun testException() = runTest {
         val flow = flow {
             emit(314)