Fix flaky shared flow subscription

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

@@ -109,14 +109,23 @@ public abstract class ChannelFlow<T>(
      * For non-atomic start it is possible to observe the situation,
      * where the pipeline after the [flowOn] call successfully executes (mostly, its `onCompletion`)
      * handlers, while the pipeline before does not, because it was cancelled during its dispatch.
-     * Thus `onCompletion` and `finally` blocks won't be executed and it may lead to a different kinds of memory leaks.
+     * Thus `onCompletion` and `finally` blocks won't be executed, and it may lead to a different kind of memory leaks.
      */
     public open fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> =
-        scope.produce(context, produceCapacity, onBufferOverflow, start = CoroutineStart.ATOMIC, block = collectToFun)
+        produceImplInternal(scope, CoroutineStart.ATOMIC)
+
+    internal open fun produceImplInternal(scope: CoroutineScope, start: CoroutineStart): ReceiveChannel<T> = scope.produce(context, produceCapacity, onBufferOverflow, start = start, block = collectToFun)
 
     override suspend fun collect(collector: FlowCollector<T>): Unit =
         coroutineScope {
-            collector.emitAll(produceImpl(this))
+            // If upstream and collect have the same dispatcher, launch the `produce` coroutine undispatched.
+            // This allows the collector to reliably subscribe to the flow before it starts emitting.
+            val current = currentCoroutineContext()[ContinuationInterceptor]
+            val desired = context[ContinuationInterceptor]
+            val start = if (desired == null || desired == current) {
+                CoroutineStart.UNDISPATCHED
+            } else CoroutineStart.ATOMIC
+            collector.emitAll(produceImplInternal(this, start))
         }
 
     protected open fun additionalToStringProps(): String? = null

kotlinx-coroutines-core/common/test/flow/channels/ChannelFlowTest.kt

@@ -206,13 +206,11 @@ class ChannelFlowTest : TestBase() {
     }
 
     @Test
-    fun testDoesntDispatchWhenUnnecessarilyWhenCollected() = runTest {
+    fun testDoesntDispatchUnnecessarilyWhenCollected() = runTest {
         expect(1)
-        var subscribed = false
         val myFlow = flow<Int> {
             expect(3)
-            subscribed = true
-            yield()
+            yield() // In other words, testing that this will be the first suspension point in `collectLatest`.
             expect(5)
         }
         launch(start = CoroutineStart.UNDISPATCHED) {
@@ -223,6 +221,20 @@ class ChannelFlowTest : TestBase() {
             finish(6)
         }
         expect(4)
-        assertTrue(subscribed)
+    }
+
+    @Test
+    fun testDispatchesToDifferentDispatcherWhenCollected() = runTest {
+        expect(1)
+        val myFlow = flow<Int> {
+            finish(4)
+        }.flowOn(wrapperDispatcher())
+        launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(2)
+            myFlow.collectLatest {
+                expectUnreached()
+            }
+        }
+        expect(3)
     }
 }