Declarative flow operators (#1291)

* Flow.onCompletion operator
* Flow.launchIn operator

Fixes #1263

Author: qwwdfsad

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -860,8 +860,10 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static synthetic fun flowWith$default (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun fold (Lkotlinx/coroutines/flow/Flow;Ljava/lang/Object;Lkotlin/jvm/functions/Function3;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun getDEFAULT_CONCURRENCY ()I
+	public static final fun launchIn (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/CoroutineScope;)Lkotlinx/coroutines/Job;
 	public static final fun map (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun mapNotNull (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun onCompletion (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun onEach (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun onErrorCollect (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;)Lkotlinx/coroutines/flow/Flow;
 	public static synthetic fun onErrorCollect$default (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
@@ -910,8 +912,8 @@ public final class kotlinx/coroutines/flow/MigrationKt {
 	public static final fun scanFold (Lkotlinx/coroutines/flow/Flow;Ljava/lang/Object;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun skip (Lkotlinx/coroutines/flow/Flow;I)Lkotlinx/coroutines/flow/Flow;
 	public static final fun subscribe (Lkotlinx/coroutines/flow/Flow;)V
-	public static final fun subscribe (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;)V
-	public static final fun subscribe (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)V
+	public static final fun subscribe (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)V
+	public static final fun subscribe (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V
 	public static final fun subscribeOn (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun withContext (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function1;)V
 }

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

@@ -18,7 +18,8 @@ import kotlin.coroutines.*
  * They only set up a chain of operations for future execution and quickly return.
  * This is known as a _cold flow_ property.
  *
- * _Terminal operators_ on the flow are suspending functions such as [collect], [single], [reduce], [toList], etc.
+ * _Terminal operators_ on the flow are either suspending functions such as [collect], [single], [reduce], [toList], etc.
+ * or [launchIn] operator that starts collection of the flow in the given scope.
  * They are applied to the upstream flow and trigger execution of all operations.
  * Execution of the flow is also called _collecting the flow_  and is always performed in a suspending manner
  * without actual blocking. Terminal operator complete normally or exceptionally depending on successful or failed
@@ -142,6 +143,7 @@ import kotlin.coroutines.*
  *     .map { computeTwo(it) }
  *     .collect { process(it) } // throws exceptions from process and computeTwo
  * ```
+ * The same reasoning can be applied to [onCompletion] operator that is a declarative replacement for `finally` block.
  *
  * Failure to adhere to the exception transparency requirement would result in strange behaviours that would make
  * it hard to reason about the code because an exception in the `collect { ... }` could be somehow "caught"

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

@@ -127,7 +127,7 @@ public fun <T, R> FlowCollector<T>.withContext(context: CoroutineContext, block:
 
 /**
  * `subscribe` is Rx-specific API that has no direct match in flows.
- * One can use `launch` instead, for example the following:
+ * One can use [launchIn] instead, for example the following:
  * ```
  * flowable
  *     .observeOn(Schedulers.io())
@@ -136,30 +136,36 @@ public fun <T, R> FlowCollector<T>.withContext(context: CoroutineContext, block:
  *
  * has the following Flow equivalent:
  * ```
- * launch(Dispatchers.IO) {
- *     try {
- *         flow.collect { value ->
- *             println("Received $value")
- *         }
- *         println("Flow is completed successfully")
- *     } catch (e: Throwable) {
- *         println("Exception $e happened")
- *     }
- * }
+ * flow
+ *     .onEach { value -> println("Received $value") }
+ *     .onCompletion { cause -> if (cause == null) println("Flow is completed successfully") }
+ *     .catch { cause -> println("Exception $cause happened") }
+ *     .flowOn(Dispatchers.IO)
+ *     .launchIn(myScope)
  * ```
- * But most of the time it is better to use terminal operators like [single] instead of [collect].
+ *
+ * Note that resulting value of [launchIn] is not used because the provided scope takes care of cancellation.
+ *
+ * Or terminal operators like [single] can be used from suspend functions.
  * @suppress
  */
-@Deprecated(message = "Use launch + collect instead", level = DeprecationLevel.ERROR)
+@Deprecated(
+    message = "Use launchIn with onEach, onCompletion and catch operators instead",
+    level = DeprecationLevel.ERROR
+)
 public fun <T> Flow<T>.subscribe(): Unit = error("Should not be called")
 
 /** @suppress **/
-@Deprecated(message = "Use launch + collect instead", level = DeprecationLevel.ERROR)
-public fun <T> Flow<T>.subscribe(onEach: (T) -> Unit): Unit = error("Should not be called")
+@Deprecated(
+    message = "Use launchIn with onEach, onCompletion and catch operators instead",
+    level = DeprecationLevel.ERROR
+)public fun <T> Flow<T>.subscribe(onEach: suspend (T) -> Unit): Unit = error("Should not be called")
 
 /** @suppress **/
-@Deprecated(message = "Use launch + collect instead", level = DeprecationLevel.ERROR)
-public fun <T> Flow<T>.subscribe(onEach: (T) -> Unit, onError: (Throwable) -> Unit): Unit = error("Should not be called")
+@Deprecated(
+    message = "Use launchIn with onEach, onCompletion and catch operators instead",
+    level = DeprecationLevel.ERROR
+)public fun <T> Flow<T>.subscribe(onEach: suspend (T) -> Unit, onError: suspend (Throwable) -> Unit): Unit = error("Should not be called")
 
 /**
  * Note that this replacement is sequential (`concat`) by default.

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

@@ -57,9 +57,8 @@ import kotlinx.coroutines.flow.unsafeFlow as flow
 @ExperimentalCoroutinesApi // tentatively stable in 1.3.0
 public fun <T> Flow<T>.catch(action: suspend FlowCollector<T>.(cause: Throwable) -> Unit): Flow<T> =
     flow {
-        catchImpl(this) { e ->
-            action(e)
-        }
+        val exception = catchImpl(this)
+        if (exception != null) action(exception)
     }
 
 /**
@@ -177,22 +176,22 @@ public fun <T> Flow<T>.retryWhen(predicate: suspend FlowCollector<T>.(cause: Thr
         var shallRetry: Boolean
         do {
             shallRetry = false
-            catchImpl(this) { e ->
-                if (predicate(e, attempt)) {
+            val cause = catchImpl(this)
+            if (cause != null) {
+                if (predicate(cause, attempt)) {
                     shallRetry = true
                     attempt++
                 } else {
-                    throw e
+                    throw cause
                 }
             }
         } while (shallRetry)
     }
 
-// Note that exception may come from the downstream operators, we should not catch them
-private suspend inline fun <T> Flow<T>.catchImpl(
-    collector: FlowCollector<T>,
-    action: FlowCollector<T>.(Throwable) -> Unit
-) {
+// Return exception from upstream or null
+internal suspend fun <T> Flow<T>.catchImpl(
+    collector: FlowCollector<T>
+): Throwable? {
     var fromDownstream: Throwable? = null
     try {
         collect {
@@ -209,10 +208,12 @@ private suspend inline fun <T> Flow<T>.catchImpl(
          * Seconds check ignores cancellation causes, they cannot be caught.
          */
         if (e.isSameExceptionAs(fromDownstream) || e.isCancellationCause(coroutineContext)) {
-            throw e
+            throw e // Rethrow exceptions from downstream and cancellation causes
+        } else {
+            return e // not from downstream
         }
-        collector.action(e)
     }
+    return null
 }
 
 private fun Throwable.isCancellationCause(coroutineContext: CoroutineContext): Boolean {

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

@@ -9,7 +9,8 @@
 package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.flow.internal.NULL
+import kotlinx.coroutines.flow.internal.*
+import kotlin.coroutines.*
 import kotlin.jvm.*
 import kotlinx.coroutines.flow.unsafeFlow as flow
 
@@ -100,6 +101,64 @@ public fun <T> Flow<T>.onEach(action: suspend (T) -> Unit): Flow<T> = flow {
     }
 }
 
+/**
+ * Invokes the given [action] when the given flow is completed or cancelled, using
+ * the exception from the upstream (if any) as cause parameter of [action].
+ *
+ * Conceptually, [onCompletion] is similar to wrapping the flow collection into a `finally` block,
+ * for example the following imperative snippet:
+ * ```
+ * try {
+ *     myFlow.collect { value ->
+ *         println(value)
+ *     }
+ * } finally {
+ *     println("Done")
+ * }
+ * ```
+ *
+ * can be replaced with a declarative one using [onCompletion]:
+ * ```
+ * myFlow
+ *     .onEach { println(it) }
+ *     .onCompletion { println("Done") }
+ *     .collect()
+ * ```
+ *
+ * This operator is *transparent* to exceptions that occur in downstream flow
+ * and does not observe exceptions that are thrown to cancel the flow,
+ * while any exception from the [action] will be thrown downstream.
+ * This behaviour can be demonstrated by the following example:
+ * ```
+ * flow { emitData() }
+ *     .map { computeOne(it) }
+ *     .onCompletion { println(it) } // Can print exceptions from emitData and computeOne
+ *     .map { computeTwo(it) }
+ *     .onCompletion { println(it) } // Can print exceptions from emitData, computeOne, onCompletion and computeTwo
+ *     .collect()
+ * ```
+ */
+@ExperimentalCoroutinesApi // tentatively stable in 1.3.0
+public fun <T> Flow<T>.onCompletion(action: suspend (cause: Throwable?) -> Unit): Flow<T> = flow {
+    var exception: Throwable? = null
+    try {
+        exception = catchImpl(this)
+    } finally {
+        // Separate method because of KT-32220
+        invokeSafely(action, exception)
+        exception?.let { throw it }
+    }
+}
+
+private suspend fun invokeSafely(action: suspend (cause: Throwable?) -> Unit, cause: Throwable?) {
+    try {
+        action(cause)
+    } catch (e: Throwable) {
+        if (cause !== null) e.addSuppressedThrowable(cause)
+        throw e
+    }
+}
+
 /**
  * Folds the given flow with [operation], emitting every intermediate result, including [initial] value.
  * Note that initial value should be immutable (or should not be mutated) as it is shared between different collectors.

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

@@ -17,7 +17,7 @@ import kotlin.jvm.*
  *
  * It is a shorthand for `collect {}`.
  *
- * This operator is usually used with [onEach] and [catch] operators to process all emitted values and
+ * This operator is usually used with [onEach], [onCompletion] and [catch] operators to process all emitted values and
  * handle an exception that might occur in the upstream flow or during processing, for example:
  *
  * ```
@@ -30,6 +30,27 @@ import kotlin.jvm.*
 @ExperimentalCoroutinesApi // tentatively stable in 1.3.0
 public suspend fun Flow<*>.collect() = collect(NopCollector)
 
+/**
+ * Terminal flow operator that [launches][launch] the [collection][collect] of the given flow in the [scope].
+ * It is a shorthand for `scope.launch { flow.collect() }`.
+ *
+ * This operator is usually used with [onEach], [onCompletion] and [catch] operators to process all emitted values
+ * handle an exception that might occur in the upstream flow or during processing, for example:
+ * ```
+ * flow
+ *     .onEach { value -> updateUi(value) }
+ *     .onCompletion { cause -> updateUi(if (cause == null) "Done" else "Failed") }
+ *     .catch { cause -> LOG.error("Exception: $cause") }
+ *     .launchIn(uiScope)
+ * ```
+ *
+ * Note that resulting value of [launchIn] is not used the provided scope takes care of cancellation.
+ */
+@ExperimentalCoroutinesApi // tentatively stable in 1.3.0
+public fun <T> Flow<T>.launchIn(scope: CoroutineScope): Job = scope.launch {
+    collect() // tail-call
+}
+
 /**
  * Terminal flow operator that collects the given flow with a provided [action].
  * If any exception occurs during collect or in the provided flow, this exception is rethrown from this method.

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

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2016-2019 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 kotlin.test.*
+
+class OnCompletionTest : TestBase() {
+
+    @Test
+    fun testOnCompletion() = runTest {
+        flow {
+            expect(1)
+            emit(2)
+            expect(4)
+        }.onEach {
+            expect(2)
+        }.onCompletion {
+            assertNull(it)
+            expect(5)
+        }.onEach {
+            expect(3)
+        }.collect()
+        finish(6)
+    }
+
+    @Test
+    fun testOnCompletionWithException() = runTest {
+        flowOf(1).onEach {
+            expect(1)
+            throw TestException()
+        }.onCompletion {
+            assertTrue(it is TestException)
+            expect(2)
+        }.catch {
+            assertTrue(it is TestException)
+            expect(3)
+        }.collect()
+        finish(4)
+    }
+
+    @Test
+    fun testOnCompletionWithExceptionDownstream() = runTest {
+        flow {
+            expect(1)
+            emit(2)
+        }.onEach {
+            expect(2)
+        }.onCompletion {
+            assertNull(it)
+            expect(4)
+        }.onEach {
+            expect(3)
+            throw TestException()
+        }.catch {
+            assertTrue(it is TestException)
+            expect(5)
+        }.collect()
+        finish(6)
+    }
+
+    @Test
+    fun testMultipleOnCompletions() = runTest {
+        flowOf(1).onCompletion {
+            assertNull(it)
+            expect(2)
+        }.onEach {
+            expect(1)
+            throw TestException()
+        }.onCompletion {
+            assertTrue(it is TestException)
+            expect(3)
+        }.catch {
+            assertTrue(it is TestException)
+            expect(4)
+        }.collect()
+        finish(5)
+    }
+
+    @Test
+    fun testExceptionFromOnCompletion() = runTest {
+        flowOf(1).onEach {
+            expect(1)
+            throw TestException()
+        }.onCompletion {
+            expect(2)
+            throw TestException2()
+        }.catch {
+            assertTrue(it is TestException2)
+            expect(3)
+        }.collect()
+        finish(4)
+    }
+
+    @Test
+    fun testContextPreservation() = runTest {
+        flowOf(1).onCompletion {
+            assertEquals("OK", NamedDispatchers.name())
+            assertNull(it)
+            expect(1)
+        }.flowOn(NamedDispatchers("OK"))
+            .onEach {
+                expect(2)
+                assertEquals("default", NamedDispatchers.nameOr("default"))
+                throw TestException()
+            }
+            .catch {
+                assertTrue(it is TestException)
+                expect(3)
+            }.collect()
+        finish(4)
+    }
+}