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

Author: qwwdfsad

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

@@ -151,8 +151,9 @@ 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 ->
+            job?.ensureActive()
             semaphore.acquire() // Acquire concurrency permit
             scope.launch {
                 try {

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)
+    }
 }
 

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

@@ -73,4 +73,20 @@ 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) } }
+            .buffer(0)
+            .take(2)
+            .toList()
+        assertEquals(listOf(1, 2), result)
+    }
 }

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

@@ -261,6 +261,21 @@ class FlowOnTest : TestBase() {
         finish(3)
     }
 
+    @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)
+    }
+
     private inner class Source(private val value: Int) {
         public var contextName: String = "unknown"