Flow cancellation (Kotlin#2028)

Make Flow more cancellation friendly:

    * flow builder is cancellable by default
    * New unambiguous currentCoroutineContext top-level function
    * New Flow.cancellable() operator

Fixes Kotlin#2026

Author: qwwdfsad

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -214,6 +214,7 @@ public final class kotlinx/coroutines/CoroutineScopeKt {
 	public static synthetic fun cancel$default (Lkotlinx/coroutines/CoroutineScope;Ljava/lang/String;Ljava/lang/Throwable;ILjava/lang/Object;)V
 	public static synthetic fun cancel$default (Lkotlinx/coroutines/CoroutineScope;Ljava/util/concurrent/CancellationException;ILjava/lang/Object;)V
 	public static final fun coroutineScope (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun currentCoroutineContext (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun ensureActive (Lkotlinx/coroutines/CoroutineScope;)V
 	public static final fun isActive (Lkotlinx/coroutines/CoroutineScope;)Z
 	public static final fun plus (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/CoroutineScope;
@@ -873,6 +874,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun buffer (Lkotlinx/coroutines/flow/Flow;I)Lkotlinx/coroutines/flow/Flow;
 	public static synthetic fun buffer$default (Lkotlinx/coroutines/flow/Flow;IILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun callbackFlow (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun cancellable (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun catch (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun channelFlow (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun collect (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

kotlinx-coroutines-core/common/src/CoroutineScope.kt

@@ -239,3 +239,19 @@ public fun CoroutineScope.cancel(message: String, cause: Throwable? = null): Uni
  * ```
  */
 public fun CoroutineScope.ensureActive(): Unit = coroutineContext.ensureActive()
+
+
+/**
+ * Returns the current [CoroutineContext] retrieved by using [kotlin.coroutines.coroutineContext].
+ * This function is an alias to avoid name clash with [CoroutineScope.coroutineContext] in a receiver position:
+ *
+ * ```
+ * launch { // this: CoroutineScope
+ *     val flow = flow<Unit> {
+ *         coroutineContext // Resolves into the context of outer launch, which is incorrect, see KT-38033
+ *         currentCoroutineContext() // Retrieves actual context where the flow is collected
+ *     }
+ * }
+ * ```
+ */
+public suspend inline fun currentCoroutineContext(): CoroutineContext = coroutineContext

kotlinx-coroutines-core/common/src/flow/Builders.kt

@@ -33,6 +33,7 @@ import kotlinx.coroutines.flow.internal.unsafeFlow as flow
  *
  * fibonacci().take(100).collect { println(it) }
  * ```
+ * Emissions from [flow] builder are [cancellable] by default.
  *
  * `emit` should happen strictly in the dispatchers of the [block] in order to preserve the flow context.
  * For example, the following code will result in an [IllegalStateException]:
@@ -49,14 +50,9 @@ import kotlinx.coroutines.flow.internal.unsafeFlow as flow
 public fun <T> flow(@BuilderInference block: suspend FlowCollector<T>.() -> Unit): Flow<T> = SafeFlow(block)
 
 // Named anonymous object
-private class SafeFlow<T>(private val block: suspend FlowCollector<T>.() -> Unit) : Flow<T> {
-    override suspend fun collect(collector: FlowCollector<T>) {
-        val safeCollector = SafeCollector(collector, coroutineContext)
-        try {
-            safeCollector.block()
-        } finally {
-            safeCollector.releaseIntercepted()
-        }
+private class SafeFlow<T>(private val block: suspend FlowCollector<T>.() -> Unit) : AbstractFlow<T>() {
+    override suspend fun collectSafely(collector: FlowCollector<T>) {
+        collector.block()
     }
 }
 

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

@@ -208,6 +208,21 @@ public fun <T> Flow<T>.flowOn(context: CoroutineContext): Flow<T> {
     }
 }
 
+/**
+ * Returns a flow which checks cancellation status on each emission and throws
+ * the corresponding cancellation cause if flow collector was cancelled.
+ * Note that [flow] builder is [cancellable] by default.
+ */
+public fun <T> Flow<T>.cancellable(): Flow<T> {
+    if (this is AbstractFlow<*>) return this // Fast-path, already cancellable
+    return unsafeFlow {
+        collect {
+            currentCoroutineContext().ensureActive()
+            emit(it)
+        }
+    }
+}
+
 /**
  * The operator that changes the context where all transformations applied to the given flow within a [builder] are executed.
  * This operator is context preserving and does not affect the context of the preceding and subsequent operations.
@@ -256,7 +271,7 @@ public fun <T, R> Flow<T>.flowWith(
          * All builders are written using scoping and no global coroutines are launched, so it is safe not to provide explicit Job.
          * It is also necessary not to mess with cancellation if multiple flowWith are used.
          */
-        val originalContext = coroutineContext.minusKey(Job)
+        val originalContext = currentCoroutineContext().minusKey(Job)
         val prepared = source.flowOn(originalContext).buffer(bufferSize)
         builder(prepared).flowOn(flowContext).buffer(bufferSize).collect { value ->
             return@collect emit(value)

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

@@ -10,7 +10,6 @@ package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.internal.*
-import kotlin.coroutines.*
 import kotlin.jvm.*
 
 // ------------------ WARNING ------------------
@@ -70,7 +69,7 @@ internal inline fun <T, R> Flow<T>.unsafeTransform(
 public fun <T> Flow<T>.onStart(
     action: suspend FlowCollector<T>.() -> Unit
 ): Flow<T> = unsafeFlow { // Note: unsafe flow is used here, but safe collector is used to invoke start action
-    val safeCollector = SafeCollector<T>(this, coroutineContext)
+    val safeCollector = SafeCollector<T>(this, currentCoroutineContext())
     try {
         safeCollector.action()
     } finally {
@@ -153,7 +152,7 @@ public fun <T> Flow<T>.onCompletion(
         throw e
     }
     // Normal completion
-    SafeCollector(this, coroutineContext).invokeSafely(action, null)
+    SafeCollector(this, currentCoroutineContext()).invokeSafely(action, null)
 }
 
 /**
@@ -178,7 +177,7 @@ public fun <T> Flow<T>.onEmpty(
         emit(it)
     }
     if (isEmpty) {
-        val collector = SafeCollector(this, coroutineContext)
+        val collector = SafeCollector(this, currentCoroutineContext())
         try {
             collector.action()
         } finally {

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

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.flow
 
+import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
 /**
@@ -41,3 +42,31 @@ public fun <T> StateFlow<T>.conflate(): Flow<T> = noImpl()
     replaceWith = ReplaceWith("this")
 )
 public fun <T> StateFlow<T>.distinctUntilChanged(): Flow<T> = noImpl()
+
+//@Deprecated(
+//    message = "isActive is resolved into the extension of outer CoroutineScope which is likely to be an error." +
+//        "Use currentCoroutineContext().isActive or cancellable() operator instead " +
+//        "or specify the receiver of isActive explicitly. " +
+//        "Additionally, flow {} builder emissions are cancellable by default.",
+//    level = DeprecationLevel.WARNING, // ERROR in 1.4
+//    replaceWith = ReplaceWith("currentCoroutineContext().isActive")
+//)
+//public val FlowCollector<*>.isActive: Boolean
+//    get() = noImpl()
+//
+//@Deprecated(
+//    message = "cancel() is resolved into the extension of outer CoroutineScope which is likely to be an error." +
+//        "Use currentCoroutineContext().cancel() instead or specify the receiver of cancel() explicitly",
+//    level = DeprecationLevel.WARNING, // ERROR in 1.4
+//    replaceWith = ReplaceWith("currentCoroutineContext().cancel(cause)")
+//)
+//public fun FlowCollector<*>.cancel(cause: CancellationException? = null): Unit = noImpl()
+//
+//@Deprecated(
+//    message = "coroutineContext is resolved into the property of outer CoroutineScope which is likely to be an error." +
+//        "Use currentCoroutineContext() instead or specify the receiver of coroutineContext explicitly",
+//    level = DeprecationLevel.WARNING, // ERROR in 1.4
+//    replaceWith = ReplaceWith("currentCoroutineContext()")
+//)
+//public val FlowCollector<*>.coroutineContext: CoroutineContext
+//    get() = noImpl()

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

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow.operators
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
+import kotlin.test.*
+
+class CancellableTest : TestBase() {
+
+    @Test
+    fun testCancellable() = runTest {
+        var sum = 0
+        val flow = (0..1000).asFlow()
+            .onEach {
+                if (it != 0) currentCoroutineContext().cancel()
+                sum += it
+            }
+
+        flow.launchIn(this).join()
+        assertEquals(500500, sum)
+        
+        sum = 0
+        flow.cancellable().launchIn(this).join()
+        assertEquals(1, sum)
+    }
+
+    @Test
+    fun testFastPath() {
+        val flow = listOf(1).asFlow()
+        assertNotSame(flow, flow.cancellable())
+
+        val cancellableFlow = flow { emit(42) }
+        assertSame(cancellableFlow, cancellableFlow.cancellable())
+    }
+}

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

@@ -58,6 +58,7 @@ class FlatMapMergeTest : FlatMapMergeBaseTest() {
                 }
                 launch {
                     expect(2)
+                    yield()
                     job.cancel()
                 }
             }

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

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.flow.internal
 
+import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 
@@ -17,7 +18,8 @@ internal actual class SafeCollector<T> actual constructor(
     private var lastEmissionContext: CoroutineContext? = null
 
     override suspend fun emit(value: T) {
-        val currentContext = coroutineContext
+        val currentContext = currentCoroutineContext()
+        currentContext.ensureActive()
         if (lastEmissionContext !== currentContext) {
             checkContext(currentContext)
             lastEmissionContext = currentContext

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

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.flow.internal
 
+import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
@@ -62,6 +63,7 @@ internal actual class SafeCollector<T> actual constructor(
 
     private fun emit(uCont: Continuation<Unit>, value: T): Any? {
         val currentContext = uCont.context
+        currentContext.ensureActive()
         // This check is triggered once per flow on happy path.
         val previousContext = lastEmissionContext
         if (previousContext !== currentContext) {

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

@@ -0,0 +1,46 @@
+/*
+ * 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 kotlinx.coroutines.channels.*
+import kotlin.test.*
+
+class FlowCancellationTest : TestBase() {
+
+    @Test
+    fun testEmitIsCooperative() = runTest {
+        val latch = Channel<Unit>(1)
+        val job = flow {
+            expect(1)
+            latch.send(Unit)
+            while (true) {
+                emit(42)
+            }
+        }.launchIn(this + Dispatchers.Default)
+
+        latch.receive()
+        expect(2)
+        job.cancelAndJoin()
+        finish(3)
+    }
+
+    @Test
+    fun testIsActiveOnCurrentContext() = runTest {
+        val latch = Channel<Unit>(1)
+        val job = flow<Unit> {
+            expect(1)
+            latch.send(Unit)
+            while (currentCoroutineContext().isActive) {
+                // Do nothing
+            }
+        }.launchIn(this + Dispatchers.Default)
+
+        latch.receive()
+        expect(2)
+        job.cancelAndJoin()
+        finish(3)
+    }
+}

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

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.flow.internal
 
+import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 
@@ -17,7 +18,8 @@ internal actual class SafeCollector<T> actual constructor(
     private var lastEmissionContext: CoroutineContext? = null
 
     override suspend fun emit(value: T) {
-        val currentContext = coroutineContext
+        val currentContext = currentCoroutineContext()
+        currentContext.ensureActive()
         if (lastEmissionContext !== currentContext) {
             checkContext(currentContext)
             lastEmissionContext = currentContext

reactive/kotlinx-coroutines-reactive/src/Convert.kt

@@ -16,7 +16,7 @@ import kotlin.coroutines.*
  * @param context -- the coroutine context from which the resulting observable is going to be signalled
  */
 @Deprecated(message = "Deprecated in the favour of consumeAsFlow()",
-    level = DeprecationLevel.WARNING,
+    level = DeprecationLevel.WARNING, // Error in 1.4
     replaceWith = ReplaceWith("this.consumeAsFlow().asPublisher()"))
 public fun <T> ReceiveChannel<T>.asPublisher(context: CoroutineContext = EmptyCoroutineContext): Publisher<T> = publish(context) {
     for (t in this@asPublisher)

reactive/kotlinx-coroutines-reactive/src/ReactiveFlow.kt

@@ -85,6 +85,7 @@ private class PublisherAsFlow<T : Any>(
             var consumed = 0L
             while (true) {
                 val value = subscriber.takeNextOrNull() ?: break
+                coroutineContext.ensureActive()
                 collector.emit(value)
                 if (++consumed == requestSize) {
                     consumed = 0L

reactive/kotlinx-coroutines-reactor/test/BackpressureTest.kt

@@ -5,6 +5,7 @@
 package kotlinx.coroutines.reactor
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.reactive.*
 import org.junit.Test
@@ -36,4 +37,17 @@ class BackpressureTest : TestBase() {
             }
         finish(3)
     }
+
+    @Test
+    fun testCooperativeCancellation() = runTest {
+        val flow = Flux.fromIterable((0L..Long.MAX_VALUE)).asFlow()
+        flow.onEach { if (it > 10) currentCoroutineContext().cancel() }.launchIn(this + Dispatchers.Default).join()
+    }
+
+    @Test
+    fun testCooperativeCancellationForBuffered() = runTest(expected = { it is CancellationException }) {
+        val flow = Flux.fromIterable((0L..Long.MAX_VALUE)).asFlow()
+        val channel = flow.onEach { if (it > 10) currentCoroutineContext().cancel() }.produceIn(this + Dispatchers.Default)
+        channel.consumeEach { /* Do nothing, just consume elements */ }
+    }
 }