Fix flaky shared flow subscription

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

@@ -109,14 +109,24 @@ 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/channels/ProduceTest.kt

@@ -1,8 +1,8 @@
 package kotlinx.coroutines.channels
 
-import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.testing.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
@@ -75,7 +75,7 @@ class ProduceTest : TestBase() {
         try {
             c.receive()
             expectUnreached()
-        } catch (e: TestCancellationException) {
+        } catch (_: TestCancellationException) {
             expect(5)
         }
         yield() // to produce
@@ -196,7 +196,7 @@ class ProduceTest : TestBase() {
         assertFailsWith<IllegalStateException> { (channel as ProducerScope<*>).awaitClose() }
         callbackFlow<Unit> {
             expect(1)
-            launch {
+            launch(start = CoroutineStart.UNDISPATCHED) {
                 expect(2)
                 assertFailsWith<IllegalStateException> {
                     awaitClose { expectUnreached() }
@@ -286,7 +286,7 @@ class ProduceTest : TestBase() {
         produced.cancel()
         try {
             source.receive()
-        } catch (e: CancellationException) {
+        } catch (_: CancellationException) {
             finish(4)
         }
     }

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

@@ -1,8 +1,8 @@
 package kotlinx.coroutines.flow
 
-import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.testing.*
 import kotlin.test.*
 
 class ChannelFlowTest : TestBase() {
@@ -20,10 +20,9 @@ class ChannelFlowTest : TestBase() {
     fun testBuffer() = runTest {
         val flow = channelFlow {
             assertTrue(trySend(1).isSuccess)
-            assertTrue(trySend(2).isSuccess)
-            assertFalse(trySend(3).isSuccess)
+            assertFalse(trySend(2).isSuccess)
         }.buffer(1)
-        assertEquals(listOf(1, 2), flow.toList())
+        assertEquals(listOf(1), flow.toList())
     }
 
     @Test
@@ -34,7 +33,7 @@ class ChannelFlowTest : TestBase() {
             assertTrue(trySend(3).isSuccess)
             assertTrue(trySend(4).isSuccess)
         }.buffer(Channel.CONFLATED)
-        assertEquals(listOf(1, 4), flow.toList()) // two elements in the middle got conflated
+        assertEquals(listOf(4), flow.toList())
     }
 
     @Test
@@ -164,14 +163,13 @@ class ChannelFlowTest : TestBase() {
         val flow = channelFlow {
             // ~ callback-based API, no children
             outerScope.launch(Job()) {
-                expect(2)
                 send(1)
                 expectUnreached()
             }
             expect(1)
         }
         assertEquals(emptyList(), flow.toList())
-        finish(3)
+        finish(2)
     }
 
     @Test
@@ -206,13 +204,12 @@ class ChannelFlowTest : TestBase() {
     }
 
     @Test
-    fun testDoesntDispatchWhenUnnecessarilyWhenCollected() = runTest {
+    fun testDoesntDispatchUnnecessarilyWhenCollected() = runTest {
+        // Test that `.collectLatest` consistently subscribes to the flow when invoked on the same dispatcher as upstream.
         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 +220,21 @@ class ChannelFlowTest : TestBase() {
             finish(6)
         }
         expect(4)
-        assertTrue(subscribed)
+    }
+
+    @Test
+    fun testDispatchesToDifferentDispatcherWhenCollected() = runTest {
+        expect(1)
+        val myFlow = flow<Int> {
+            expect(4)
+        }.flowOn(wrapperDispatcher())
+        launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(2)
+            myFlow.collectLatest {
+                expectUnreached()
+            }
+            finish(5)
+        }
+        expect(3)
     }
 }

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

@@ -2,9 +2,9 @@
 
 package kotlinx.coroutines.flow
 
-import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.testing.*
 import kotlin.test.*
 
 class FlowCallbackTest : TestBase() {
@@ -14,24 +14,25 @@ class FlowCallbackTest : TestBase() {
         val flow = callbackFlow {
             // ~ callback-based API
             outerScope.launch(Job()) {
-                expect(2)
                 try {
+                    expect(4)
                     send(1)
                     expectUnreached()
                 } catch (e: IllegalStateException) {
-                    expect(3)
+                    expect(5)
                     assertTrue(e.message!!.contains("awaitClose"))
                 }
+                finish(6)
             }
             expect(1)
         }
         try {
             flow.collect()
         } catch (e: IllegalStateException) {
-            expect(4)
+            expect(2)
             assertTrue(e.message!!.contains("awaitClose"))
         }
-        finish(5)
+        expect(3)
     }
 
     @Test

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

@@ -1,14 +1,13 @@
 package kotlinx.coroutines.flow
 
-import kotlinx.coroutines.testing.*
-import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.testing.*
 import kotlin.test.*
 
 /**
  * A _behavioral_ test for conflation options that can be configured by the [buffer] operator to test that it is
  * implemented properly and that adjacent [buffer] calls are fused properly.
-*/
+ */
 class BufferConflationTest : TestBase() {
     private val n = 100 // number of elements to emit for test
 
@@ -20,8 +19,8 @@ class BufferConflationTest : TestBase() {
         expect(1)
         // emit all and conflate, then collect first & last
         val expectedList = when (onBufferOverflow) {
-            BufferOverflow.DROP_OLDEST -> listOf(0) + (n - capacity until n).toList() // first item & capacity last ones
-            BufferOverflow.DROP_LATEST -> (0..capacity).toList() // first & capacity following ones
+            BufferOverflow.DROP_OLDEST -> (n - capacity until n).toList() // first item & capacity last ones
+            BufferOverflow.DROP_LATEST -> (0 until capacity).toList() // first & capacity following ones
             else -> error("cannot happen")
         }
         flow {
@@ -38,6 +37,26 @@ class BufferConflationTest : TestBase() {
         finish(n + 2 + expectedList.size)
     }
 
+    @Test
+    fun testConflateExplicit() = runTest {
+        val n = 3
+        expect(1)
+        flow {
+            expect(2)
+            emit(0)
+            expect(3)
+            emit(1)
+            expect(4)
+            emit(2)
+        }
+            .conflate()
+            .collect { value ->
+                assertEquals(2, value)
+                expect(n + 2)
+            }
+        finish(n + 2 + 1)
+    }
+
     @Test
     fun testConflate() =
         checkConflate(1) {
@@ -138,6 +157,6 @@ class BufferConflationTest : TestBase() {
     fun testBuffer3DropOldestOverrideBuffer8DropLatest() =
         checkConflate(3, BufferOverflow.DROP_OLDEST) {
             buffer(8, onBufferOverflow = BufferOverflow.DROP_LATEST)
-            .buffer(3, BufferOverflow.DROP_OLDEST)
+                .buffer(3, BufferOverflow.DROP_OLDEST)
         }
 }