Rethrow AbortFlowException upon cancellation (Kotlin#4034)

collectWhile catches AbortFlowException even if the CoroutineScope is cancelled
while waiting on child tasks, causing it to fail to respond to cancel().

Author: jiaoxiaodong

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

@@ -5,6 +5,7 @@ package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.internal.*
+import kotlin.coroutines.*
 import kotlin.jvm.*
 import kotlinx.coroutines.flow.flow as safeFlow
 import kotlinx.coroutines.flow.internal.unsafeFlow as flow
@@ -133,5 +134,8 @@ internal suspend inline fun <T> Flow<T>.collectWhile(crossinline predicate: susp
         collect(collector)
     } catch (e: AbortFlowException) {
         e.checkOwnership(collector)
+        // If thrown in a CoroutineScope waiting on child tasks, this exception can override a later CancellationException
+        // from cancel(). Rethrow if it does.
+        if (!coroutineContext.isActive) { throw e }
     }
 }

kotlinx-coroutines-core/common/test/flow/terminal/FirstTest.kt

@@ -2,9 +2,11 @@ package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.*
+import kotlinx.coroutines.CoroutineStart.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.flow.internal.*
 import kotlin.test.*
+import kotlin.time.*
 
 class FirstTest : TestBase() {
     @Test
@@ -173,4 +175,19 @@ class FirstTest : TestBase() {
 
         assertFailsWith<CancellationException> { flow.first() }
     }
-}
+
+    @Test
+    fun test() = runTest {
+        val job = launch(start = UNDISPATCHED) {
+            flow {
+                coroutineScope {
+                    emitAll(produce(start = UNDISPATCHED) {
+                        channel.send(Unit)
+                    })
+                }
+            }.first()
+            expectUnreached()
+        }
+        job.cancel()
+    }
+}