Mark Flow.collect as internal to prevent its direct implementation and provide AbstractFlow instead that enforces context preservation guarantees

Author: qwwdfsad

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

@@ -780,6 +780,12 @@ public final class kotlinx/coroutines/channels/TickerMode : java/lang/Enum {
 	public static fun values ()[Lkotlinx/coroutines/channels/TickerMode;
 }
 
+public abstract class kotlinx/coroutines/flow/AbstractFlow : kotlinx/coroutines/flow/Flow {
+	public fun <init> ()V
+	public final fun collect (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public abstract fun collectSafely (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public abstract interface class kotlinx/coroutines/flow/Flow {
 	public abstract fun collect (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }

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

@@ -5,6 +5,8 @@
 package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.internal.SafeCollector
+import kotlin.coroutines.*
 
 /**
  * A cold asynchronous data stream that sequentially emits values
@@ -112,17 +114,59 @@ import kotlinx.coroutines.*
 @ExperimentalCoroutinesApi
 public interface Flow<out T> {
 
+    /**
+     * Accepts the given [collector] and [emits][FlowCollector.emit] values into it.
+     * This method should never be implemented or used directly.
+     *
+     * The only way to implement flow interface directly is to extend [AbstractFlow].
+     * To collect it into the specific collector, either `collector.emitAll(flow)` or `collect { }` extension should be used.
+     * Such limitation ensures that context preservation property is not violated and prevents most of the developer mistakes
+     * related to concurrency, inconsistent flow dispatchers and cancellation.
+     */
+    @InternalCoroutinesApi
+    public suspend fun collect(collector: FlowCollector<T>)
+}
+
+/**
+ * Base class to extend to have a stateful implementation of the flow.
+ * It tracks all the properties required for context preservation and throws [IllegalStateException] if any of the properties are violated.
+ * Example of the implementation:
+ * ```
+ * // list.asFlow() + collect counter
+ * class CountingListFlow(private val values: List<Int>) : AbstractFlow<Int>() {
+ *     private val collectedCounter = AtomicInteger(0)
+ *
+ *     override suspend fun collectSafely(collector: FlowCollector<Int>) {
+ *         collectedCounter.incrementAndGet() // Increment collected counter
+ *         values.forEach { // Emit all the values
+ *             collector.emit(it)
+ *         }
+ *     }
+ *
+ *     fun toDiagnosticString(): String = "Flow with values $values was collected ${collectedCounter.value} times"
+ * }
+ * ```
+ */
+@FlowPreview
+public abstract class AbstractFlow<T> : Flow<T> {
+
+    @InternalCoroutinesApi
+    public final override suspend fun collect(collector: FlowCollector<T>) {
+        collectSafely(SafeCollector(collector, collectContext = coroutineContext))
+    }
+
     /**
      * Accepts the given [collector] and [emits][FlowCollector.emit] values into it.
      *
      * A valid implementation of this method has the following constraints:
      * 1) It should not change the coroutine context (e.g. with `withContext(Dispatchers.IO)`) when emitting values.
      *    The emission should happen in the context of the [collect] call.
      *    Please refer to the top-level [Flow] documentation for more details.
-     *
      * 2) It should serialize calls to [emit][FlowCollector.emit] as [FlowCollector] implementations are not
-     * thread safe by default.
+     *    thread-safe by default.
      *    To automatically serialize emissions [channelFlow] builder can be used instead of [flow]
+     *
+     * @throws IllegalStateException if any of the invariants are violated.
      */
-    public suspend fun collect(collector: FlowCollector<T>)
+    public abstract suspend fun collectSafely(collector: FlowCollector<T>)
 }

kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt

@@ -7,12 +7,37 @@ package kotlinx.coroutines.flow
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlin.coroutines.*
+import kotlin.reflect.*
 import kotlin.test.*
 
 class FlowInvariantsTest : TestBase() {
 
+    private fun <T> runParametrizedTest(
+        expectedException: KClass<out Throwable>? = null,
+        testBody: suspend (flowFactory: (suspend FlowCollector<T>.() -> Unit) -> Flow<T>) -> Unit
+    ) = runTest {
+        val r1 = runCatching { testBody { flow(it) } }.exceptionOrNull()
+        check(r1, expectedException)
+        reset()
+
+        val r2 = runCatching { testBody { abstractFlow(it) } }.exceptionOrNull()
+        check(r2, expectedException)
+    }
+
+    private fun <T> abstractFlow(block: suspend FlowCollector<T>.() -> Unit): Flow<T> = object : AbstractFlow<T>() {
+        override suspend fun collectSafely(collector: FlowCollector<T>) {
+            collector.block()
+        }
+    }
+
+    private fun check(exception: Throwable?, expectedException: KClass<out Throwable>?) {
+        if (expectedException != null && exception == null) fail("Expected $expectedException, but test completed successfully")
+        if (expectedException != null && exception != null) assertTrue(expectedException.isInstance(exception))
+        if (expectedException == null && exception != null) throw exception
+    }
+
     @Test
-    fun testWithContextContract() = runTest({ it is IllegalStateException }) {
+    fun testWithContextContract() = runParametrizedTest<Int>(IllegalStateException::class) { flow ->
         flow {
             kotlinx.coroutines.withContext(NonCancellable) {
                 emit(1)
@@ -23,7 +48,7 @@ class FlowInvariantsTest : TestBase() {
     }
 
     @Test
-    fun testWithDispatcherContractViolated() = runTest({ it is IllegalStateException }) {
+    fun testWithDispatcherContractViolated() = runParametrizedTest<Int>(IllegalStateException::class) { flow ->
         flow {
             kotlinx.coroutines.withContext(NamedDispatchers("foo")) {
                 emit(1)
@@ -34,7 +59,7 @@ class FlowInvariantsTest : TestBase() {
     }
 
     @Test
-    fun testCachedInvariantCheckResult() = runTest {
+    fun testCachedInvariantCheckResult() = runParametrizedTest<Int> { flow ->
         flow {
             emit(1)
 
@@ -55,7 +80,7 @@ class FlowInvariantsTest : TestBase() {
     }
 
     @Test
-    fun testWithNameContractViolated() = runTest({ it is IllegalStateException }) {
+    fun testWithNameContractViolated() = runParametrizedTest<Int>(IllegalStateException::class) { flow ->
         flow {
             kotlinx.coroutines.withContext(CoroutineName("foo")) {
                 emit(1)
@@ -86,25 +111,25 @@ class FlowInvariantsTest : TestBase() {
     }
 
     @Test
-    fun testScopedJob() = runTest({ it is IllegalStateException }) {
-        flow { emit(1) }.buffer(EmptyCoroutineContext).collect {
+    fun testScopedJob() = runParametrizedTest<Int>(IllegalStateException::class) { flow ->
+        flow { emit(1) }.buffer(EmptyCoroutineContext, flow).collect {
             expect(1)
         }
 
         finish(2)
     }
 
     @Test
-    fun testScopedJobWithViolation() = runTest({ it is IllegalStateException }) {
-        flow { emit(1) }.buffer(Dispatchers.Unconfined).collect {
+    fun testScopedJobWithViolation() = runParametrizedTest<Int>(IllegalStateException::class) { flow ->
+        flow { emit(1) }.buffer(Dispatchers.Unconfined, flow).collect {
             expect(1)
         }
 
         finish(2)
     }
 
     @Test
-    fun testMergeViolation() = runTest {
+    fun testMergeViolation() = runParametrizedTest<Int> { flow ->
         fun Flow<Int>.merge(other: Flow<Int>): Flow<Int> = flow {
             coroutineScope {
                 launch {
@@ -130,17 +155,6 @@ class FlowInvariantsTest : TestBase() {
         assertFailsWith<IllegalStateException> { flow.trickyMerge(flow).toList() }
     }
 
-    // TODO merge artifact
-    private fun <T> channelFlow(bufferSize: Int = 16, @BuilderInference block: suspend ProducerScope<T>.() -> Unit): Flow<T> =
-        flow {
-            coroutineScope {
-                val channel = produce(capacity = bufferSize, block = block)
-                channel.consumeEach { value ->
-                    emit(value)
-                }
-            }
-        }
-
     @Test
     fun testNoMergeViolation() = runTest {
         fun Flow<Int>.merge(other: Flow<Int>): Flow<Int> = channelFlow {
@@ -167,7 +181,7 @@ class FlowInvariantsTest : TestBase() {
     }
 
     @Test
-    fun testScopedCoroutineNoViolation() = runTest {
+    fun testScopedCoroutineNoViolation() = runParametrizedTest<Int> { flow ->
         fun Flow<Int>.buffer(): Flow<Int> = flow {
             coroutineScope {
                 val channel = produce {
@@ -180,11 +194,10 @@ class FlowInvariantsTest : TestBase() {
                 }
             }
         }
-
         assertEquals(listOf(1, 1), flowOf(1, 1).buffer().toList())
     }
 
-    private fun Flow<Int>.buffer(coroutineContext: CoroutineContext): Flow<Int> = flow {
+    private fun Flow<Int>.buffer(coroutineContext: CoroutineContext, flow: (suspend FlowCollector<Int>.() -> Unit) -> Flow<Int>): Flow<Int> = flow {
         coroutineScope {
             val channel = Channel<Int>()
             launch {