Properly check identity of caught AbortFlowException in Flow.first operator.

It fixes two problems:
    * NoSuchElementException can be thrown during cancellation sequence (see FirstJvmTest that reproduces this problem with explanation)
    * Cancellation can be accidentally suppressed and flow activity can be prolonged

Fixes #2051

Author: qwwdfsad

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

@@ -84,15 +84,18 @@ public suspend fun <T: Any> Flow<T>.singleOrNull(): T? {
  */
 public suspend fun <T> Flow<T>.first(): T {
     var result: Any? = NULL
-    try {
-        collect { value ->
+    val collector = object : FlowCollector<T> {
+        override suspend fun emit(value: T) {
             result = value
-            throw AbortFlowException(NopCollector)
+            throw AbortFlowException(this)
         }
+    }
+    try {
+        collect(collector)
     } catch (e: AbortFlowException) {
-        // Do nothing
+        // Do not suppress other cancellation sources
+        e.checkOwnership(collector)
     }
-
     if (result === NULL) throw NoSuchElementException("Expected at least one element")
     return result as T
 }
@@ -103,17 +106,20 @@ public suspend fun <T> Flow<T>.first(): T {
  */
 public suspend fun <T> Flow<T>.first(predicate: suspend (T) -> Boolean): T {
     var result: Any? = NULL
-    try {
-        collect { value ->
+    val collector = object : FlowCollector<T> {
+        override suspend fun emit(value: T) {
             if (predicate(value)) {
                 result = value
-                throw AbortFlowException(NopCollector)
+                throw AbortFlowException(this)
             }
         }
+    }
+    try {
+        collect(collector)
     } catch (e: AbortFlowException) {
-        // Do nothing
+        // Do not suppress other cancellation sources
+        e.checkOwnership(collector)
     }
-
     if (result === NULL) throw NoSuchElementException("Expected at least one element matching the predicate $predicate")
     return result as T
 }
@@ -124,13 +130,17 @@ public suspend fun <T> Flow<T>.first(predicate: suspend (T) -> Boolean): T {
  */
 public suspend fun <T : Any> Flow<T>.firstOrNull(): T? {
     var result: T? = null
-    try {
-        collect { value ->
+    val collector = object : FlowCollector<T> {
+        override suspend fun emit(value: T) {
             result = value
-            throw AbortFlowException(NopCollector)
+            throw AbortFlowException(this)
         }
+    }
+    try {
+        collect(collector)
     } catch (e: AbortFlowException) {
-        // Do nothing
+        // Do not suppress other cancellation sources
+        e.checkOwnership(collector)
     }
     return result
 }
@@ -141,15 +151,19 @@ public suspend fun <T : Any> Flow<T>.firstOrNull(): T? {
  */
 public suspend fun <T : Any> Flow<T>.firstOrNull(predicate: suspend (T) -> Boolean): T? {
     var result: T? = null
-    try {
-        collect { value ->
+    val collector = object : FlowCollector<T> {
+        override suspend fun emit(value: T) {
             if (predicate(value)) {
                 result = value
-                throw AbortFlowException(NopCollector)
+                throw AbortFlowException(this)
             }
         }
+    }
+    try {
+        collect(collector)
     } catch (e: AbortFlowException) {
-        // Do nothing
+        // Do not suppress other cancellation sources
+        e.checkOwnership(collector)
     }
     return result
 }

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

@@ -6,6 +6,7 @@ package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.internal.*
 import kotlin.test.*
 
 class FirstTest : TestBase() {
@@ -160,4 +161,13 @@ class FirstTest : TestBase() {
         assertSame(instance, flow.first { true })
         assertSame(instance, flow.firstOrNull { true })
     }
+
+    @Test
+    fun testAbortFlowException() = runTest {
+        val flow = flow<Int> {
+            throw AbortFlowException(NopCollector) // Emulate cancellation
+        }
+
+        assertFailsWith<CancellationException> { flow.first() }
+    }
 }

kotlinx-coroutines-core/jvm/test/flow/FirstJvmTest.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import org.junit.Test
+import kotlin.test.*
+
+class FirstJvmTest : TestBase() {
+
+    @Test
+    fun testTakeInterference() = runBlocking(Dispatchers.Default) {
+        /*
+         * This test tests a racy situation when outer channelFlow is being cancelled,
+         * inner flow starts atomically in "CANCELLING" state, sends one element and completes
+         * (=> cancels and drops element away), triggering NSEE in Flow.first operator
+         */
+        val values = (0..10000).asFlow().flatMapMerge(Int.MAX_VALUE) {
+            channelFlow {
+                val value = channelFlow { send(1) }.first()
+                send(value)
+            }
+        }.take(1).toList()
+        assertEquals(listOf(1), values)
+    }
+}