Strengthen flow context preservation invariant

    * Add additional check in SafeCollector with an error message pointing to channelFlow
    * Improve performance of the CoroutineId check in SafeCollector

Fixes #1210

Author: qwwdfsad

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

@@ -27,7 +27,7 @@ import kotlinx.coroutines.*
  * trigger their evaluation every time [collect] is executed) or hot ones, but, conventionally, they represent cold streams.
  * Transitions between hot and cold streams are supported via channels and the corresponding API: [flowViaChannel], [broadcastIn], [produceIn].
  *
- * The flow has a context preserving property: it encapsulates its own execution context and never propagates or leaks it downstream, thus making
+ * The flow has a context preservation property: it encapsulates its own execution context and never propagates or leaks it downstream, thus making
  * reasoning about the execution context of particular transformations or terminal operations trivial.
  *
  * There are two ways to change the context of a flow: [flowOn][Flow.flowOn] and [flowWith][Flow.flowWith].
@@ -104,6 +104,7 @@ public interface Flow<out T> {
      * is a proper [Flow] implementation, but using `launch(Dispatchers.IO)` is not.
      *
      * 2) It should serialize calls to [emit][FlowCollector.emit] as [FlowCollector] implementations are not thread safe by default.
+     *    To automatically serialize emissions [channelFlow] builder can be used instead of [flow]
      */
     public suspend fun collect(collector: FlowCollector<T>)
 }

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

@@ -6,36 +6,50 @@ package kotlinx.coroutines.flow.internal
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
-import kotlinx.coroutines.internal.*
+import kotlinx.coroutines.internal.ScopeCoroutine
 import kotlin.coroutines.*
 
 @PublishedApi
 internal class SafeCollector<T>(
     private val collector: FlowCollector<T>,
     collectContext: CoroutineContext
-) : FlowCollector<T>, SynchronizedObject() {
+) : FlowCollector<T> {
 
     private val collectContext = collectContext.minusKey(Job).minusId()
-    private var lastObservedContext: CoroutineContext? = null
+    private var lastEmissionContext: CoroutineContext? = null
 
     override suspend fun emit(value: T)  {
         /*
          * Benign data-race here:
          * We read potentially racy published coroutineContext, but we only use it for
-         * referential comparison (=> thus safe) and are not using it for actual comparisons.
+         * referential comparison (=> thus safe) and are not using it for structural comparisons.
          */
         val currentContext = coroutineContext
-        if (lastObservedContext !== currentContext) {
+        val observedContext = lastEmissionContext
+        if (observedContext !== currentContext) {
+            if (observedContext !== null) checkJobs(observedContext, currentContext)
             val emitContext = currentContext.minusKey(Job).minusId()
             if (emitContext != collectContext) {
-                error(
-                    "Flow invariant is violated: flow was collected in $collectContext, but emission happened in $emitContext. " +
-                            "Please refer to 'flow' documentation or use 'flowOn' instead"
-                )
-            }
-            // Racy publication
-            lastObservedContext = currentContext
+                    error(
+                        "Flow invariant is violated: flow was collected in $collectContext, but emission happened in $emitContext. " +
+                                "Please refer to 'flow' documentation or use 'flowOn' instead"
+                    )
+                }
+            lastEmissionContext = currentContext
         }
         collector.emit(value) // TCE
     }
+
+    private fun checkJobs(observedContext: CoroutineContext, currentContext: CoroutineContext) {
+        val previousJob = observedContext[Job].transitiveCoroutineParent()
+        val currentJob = currentContext[Job].transitiveCoroutineParent()
+        check(previousJob === currentJob) { "Flow invariant is violated: emissions from different coroutines are detected ($currentContext and $lastEmissionContext). " +
+                "FlowCollector is not thread-safe and concurrent emissions are prohibited. To mitigate this restriction please use 'flowChannel' builder instead of 'flow'" }
+    }
+
+    private fun Job?.transitiveCoroutineParent(): Job? {
+        if (this === null) return null
+        if (this !is ScopeCoroutine<*>) return this
+        return parent.transitiveCoroutineParent()
+    }
 }

kotlinx-coroutines-core/common/src/internal/Scopes.kt

@@ -19,6 +19,8 @@ internal open class ScopeCoroutine<in T>(
     final override fun getStackTraceElement(): StackTraceElement? = null
     override val defaultResumeMode: Int get() = MODE_DIRECT
 
+    internal val parent: Job? get() = parentContext[Job]
+
     override val cancelsParent: Boolean
         get() = false // it throws exception to parent instead of cancelling it
 

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

@@ -104,6 +104,71 @@ class FlowInvariantsTest : TestBase() {
         finish(2)
     }
 
+    @Test
+    fun testMergeViolation() = runTest {
+        fun Flow<Int>.merge(other: Flow<Int>): Flow<Int> = flow {
+            coroutineScope {
+                launch {
+                    collect { value -> emit(value) }
+                }
+                other.collect { value -> emit(value) }
+            }
+        }
+
+        fun Flow<Int>.trickyMerge(other: Flow<Int>): Flow<Int> = flow {
+            coroutineScope {
+                launch {
+                    collect { value ->
+                        coroutineScope { emit(value) }
+                    }
+                }
+                other.collect { value -> emit(value) }
+            }
+        }
+
+        val flow = flowOf(1)
+        assertFailsWith<IllegalStateException> { flow.merge(flow).toList() }
+        assertFailsWith<IllegalStateException> { flow.trickyMerge(flow).toList() }
+    }
+
+
+    // TODO merge artifact
+    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 {
+            launch {
+                collect { value -> send(value) }
+            }
+            other.collect { value -> send(value) }
+        }
+
+        fun Flow<Int>.trickyMerge(other: Flow<Int>): Flow<Int> = channelFlow {
+            coroutineScope {
+                launch {
+                    collect { value ->
+                        coroutineScope { send(value) }
+                    }
+                }
+                other.collect { value -> send(value) }
+            }
+        }
+
+        val flow = flowOf(1)
+        assertEquals(listOf(1, 1), flow.merge(flow).toList())
+        assertEquals(listOf(1, 1), flow.trickyMerge(flow).toList())
+    }
+
+
     private fun Flow<Int>.buffer(coroutineContext: CoroutineContext): Flow<Int> = flow {
         coroutineScope {
             val channel = Channel<Int>()

kotlinx-coroutines-core/jvm/src/CoroutineContext.kt

@@ -80,7 +80,7 @@ internal actual val CoroutineContext.coroutineName: String? get() {
 }
 
 @Suppress("NOTHING_TO_INLINE")
-internal actual inline fun CoroutineContext.minusId(): CoroutineContext = minusKey(CoroutineId)
+internal actual inline fun CoroutineContext.minusId(): CoroutineContext =  if (DEBUG) minusKey(CoroutineId) else this
 
 private const val DEBUG_THREAD_NAME_SEPARATOR = " @"