Properly distinguish AbortFlowExceptions from different non-terminal operators

Fixes #1610

Author: qwwdfsad

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

@@ -114,7 +114,7 @@ internal fun <T1, T2, R> zipImpl(flow: Flow<T1>, flow2: Flow<T2>, transform: sus
          * Invariant: this clause is invoked only when all elements from the channel were processed (=> rendezvous restriction).
          */
         (second as SendChannel<*>).invokeOnClose {
-            if (!first.isClosedForReceive) first.cancel(AbortFlowException())
+            if (!first.isClosedForReceive) first.cancel(AbortFlowException(this@unsafeFlow))
         }
 
         val otherIterator = second.iterator()
@@ -126,9 +126,9 @@ internal fun <T1, T2, R> zipImpl(flow: Flow<T1>, flow2: Flow<T2>, transform: sus
                 emit(transform(NULL.unbox(value), NULL.unbox(otherIterator.next())))
             }
         } catch (e: AbortFlowException) {
-            // complete
+            e.checkOwnership(owner = this@unsafeFlow)
         } finally {
-            if (!second.isClosedForReceive) second.cancel(AbortFlowException())
+            if (!second.isClosedForReceive) second.cancel(AbortFlowException(this@unsafeFlow))
         }
     }
 }

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

@@ -5,12 +5,21 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 
 /**
  * This exception is thrown when operator need no more elements from the flow.
  * This exception should never escape outside of operator's implementation.
+ * This exception can be safely ignored by non-terminal flow operator if and only if it was caught by its owner
+ * (see usages of [checkOwnership]).
  */
-internal expect class AbortFlowException() : CancellationException
+internal expect class AbortFlowException(owner: FlowCollector<*>) : CancellationException {
+    public val owner: FlowCollector<*>
+}
+
+internal fun AbortFlowException.checkOwnership(owner: FlowCollector<*>) {
+    if (this.owner !== owner) throw this
+}
 
 /**
  * Exception used to cancel child of [scopedFlow] without cancelling the whole scope.

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

@@ -62,14 +62,14 @@ public fun <T> Flow<T>.take(count: Int): Flow<T> {
                 }
             }
         } catch (e: AbortFlowException) {
-            // Nothing, bail out
+            e.checkOwnership(owner = this)
         }
     }
 }
 
 private suspend fun <T> FlowCollector<T>.emitAbort(value: T) {
     emit(value)
-    throw AbortFlowException()
+    throw AbortFlowException(this)
 }
 
 /**
@@ -80,9 +80,9 @@ public fun <T> Flow<T>.takeWhile(predicate: suspend (T) -> Boolean): Flow<T> = f
     try {
         collect { value ->
             if (predicate(value)) emit(value)
-            else throw AbortFlowException()
+            else throw AbortFlowException(this)
         }
     } catch (e: AbortFlowException) {
-        // Nothing, bail out
+        e.checkOwnership(owner = this)
     }
 }

kotlinx-coroutines-core/common/src/flow/terminal/Collect.kt

@@ -130,5 +130,6 @@ public suspend fun <T> Flow<T>.collectLatest(action: suspend (value: T) -> Unit)
  * Collects all the values from the given [flow] and emits them to the collector.
  * It is a shorthand for `flow.collect { value -> emit(value) }`.
  */
+@BuilderInference
 @ExperimentalCoroutinesApi
 public suspend inline fun <T> FlowCollector<T>.emitAll(flow: Flow<T>) = flow.collect(this)

kotlinx-coroutines-core/common/src/flow/terminal/Reduce.kt

@@ -90,7 +90,7 @@ public suspend fun <T> Flow<T>.first(): T {
     try {
         collect { value ->
             result = value
-            throw AbortFlowException()
+            throw AbortFlowException(NopCollector)
         }
     } catch (e: AbortFlowException) {
         // Do nothing
@@ -110,7 +110,7 @@ public suspend fun <T> Flow<T>.first(predicate: suspend (T) -> Boolean): T {
         collect { value ->
             if (predicate(value)) {
                 result = value
-                throw AbortFlowException()
+                throw AbortFlowException(NopCollector)
             }
         }
     } catch (e: AbortFlowException) {

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

@@ -131,4 +131,18 @@ class TakeTest : TestBase() {
             }
         finish(2)
     }
+
+    @Test
+    fun testNestedTake() = runTest {
+        val inner = flow {
+            emit(1)
+            expectUnreached()
+        }.take(1)
+        val outer = flow {
+            while(true) {
+                emitAll(inner)
+            }
+        }
+        assertEquals(listOf(1, 1, 1), outer.take(3).toList())
+    }
 }

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

@@ -44,7 +44,7 @@ class TransformLatestTest : TestBase() {
     }
 
     @Test
-    fun testSwitchRendevouzBuffer() = runTest {
+    fun testSwitchRendezvousBuffer() = runTest {
         val flow = flowOf(1, 2, 3, 4, 5)
         flow.transformLatest {
             emit(it)
@@ -157,16 +157,4 @@ class TransformLatestTest : TestBase() {
         val flow = flowOf(1, 2, 3, 4, 5).transformLatest { emit(it) }
         assertEquals(listOf(1), flow.take(1).toList())
     }
-
-    @Test
-    @Ignore // TODO separate branch and/or discuss
-    fun testTakeUpstreamCancellation() = runTest {
-        val flow = flow {
-            emit(1)
-            expectUnreached()
-            emit(2)
-            emit(3)
-        }.transformLatest { emit(it) }
-        assertEquals(listOf(1), flow.take(1).toList())
-    }
 }

kotlinx-coroutines-core/js/src/flow/internal/FlowExceptions.kt

@@ -5,6 +5,9 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 
-internal actual class AbortFlowException : CancellationException("Flow was aborted, no more elements needed")
+internal actual class AbortFlowException actual constructor(
+    actual val owner: FlowCollector<*>
+) : CancellationException("Flow was aborted, no more elements needed")
 internal actual class ChildCancelledException : CancellationException("Child of the scoped flow was cancelled")

kotlinx-coroutines-core/jvm/src/flow/internal/FlowExceptions.kt

@@ -5,8 +5,12 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
+
+internal actual class AbortFlowException actual constructor(
+    actual val owner: FlowCollector<*>
+) : CancellationException("Flow was aborted, no more elements needed") {
 
-internal actual class AbortFlowException : CancellationException("Flow was aborted, no more elements needed") {
     override fun fillInStackTrace(): Throwable {
         if (DEBUG) super.fillInStackTrace()
         return this

kotlinx-coroutines-core/native/src/flow/internal/FlowExceptions.kt

@@ -5,7 +5,10 @@
 package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 
-internal actual class AbortFlowException : CancellationException("Flow was aborted, no more elements needed")
+internal actual class AbortFlowException actual constructor(
+    actual val owner: FlowCollector<*>
+) : CancellationException("Flow was aborted, no more elements needed")
 internal actual class ChildCancelledException : CancellationException("Child of the scoped flow was cancelled")