Fix flowOn handling of thread context elements (#4431)

dkhalanskyjb · web-flow · commit 307d65b72eb7 · 2025-05-21T11:09:45.000+02:00
Fixes #4403 Fixes #4422 Fixes some other, similar bugs that weren't reported.
diff --git a/kotlinx-coroutines-core/common/src/channels/Produce.kt b/kotlinx-coroutines-core/common/src/channels/Produce.kt
@@ -266,6 +266,7 @@ public fun <E> CoroutineScope.produce(
     produce(context, capacity, BufferOverflow.SUSPEND, start, onCompletion, block)
 
 // Internal version of produce that is maximally flexible, but is not exposed through public API (too many params)
+// (scope + context1).produce(context2) == scope.produce(context1 + context2)
 internal fun <E> CoroutineScope.produce(
     context: CoroutineContext = EmptyCoroutineContext,
     capacity: Int = 0,
diff --git a/kotlinx-coroutines-core/common/src/flow/internal/ChannelFlow.kt b/kotlinx-coroutines-core/common/src/flow/internal/ChannelFlow.kt
@@ -217,12 +217,11 @@ internal suspend fun <T, V> withContextUndispatched(
     value: V,
     countOrElement: Any = threadContextElements(newContext), // can be precomputed for speed
     block: suspend (V) -> T
-): T =
+): T = withCoroutineContext(newContext, countOrElement) {
     suspendCoroutineUninterceptedOrReturn { uCont ->
-        withCoroutineContext(newContext, countOrElement) {
-            block.startCoroutineUninterceptedOrReturn(value, StackFrameContinuation(uCont, newContext))
-        }
+        block.startCoroutineUninterceptedOrReturn(value, StackFrameContinuation(uCont, newContext))
     }
+}
 
 // Continuation that links the caller with uCont with walkable CoroutineStackFrame
 private class StackFrameContinuation<T>(
diff --git a/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt b/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt
@@ -9,6 +9,13 @@ import kotlin.coroutines.jvm.internal.CoroutineStackFrame
  * [ContinuationInterceptor] is specified and adds optional support for debugging facilities (when turned on)
  * and copyable-thread-local facilities on JVM.
  * See [DEBUG_PROPERTY_NAME] for description of debugging facilities on JVM.
+ *
+ * When [CopyableThreadContextElement] values are used, the logic for processing them is as follows:
+ * - If [this] or [context] has a copyable thread-local value whose key is absent in [context] or [this],
+ *   it is [copied][CopyableThreadContextElement.copyForChild] to the new context.
+ * - If [this] has a copyable thread-local value whose key is present in [context],
+ *   it is [merged][CopyableThreadContextElement.mergeForChild] with the one from [context].
+ * - The other values are added to the new context as is, with [context] values taking precedence.
  */
 @ExperimentalCoroutinesApi
 public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
@@ -37,13 +44,14 @@ private fun CoroutineContext.hasCopyableElements(): Boolean =
 
 /**
  * Folds two contexts properly applying [CopyableThreadContextElement] rules when necessary.
- * The rules are the following:
- * - If neither context has CTCE, the sum of two contexts is returned
- * - Every CTCE from the left-hand side context that does not have a matching (by key) element from right-hand side context
- *   is [copied][CopyableThreadContextElement.copyForChild] if [isNewCoroutine] is `true`.
- * - Every CTCE from the left-hand side context that has a matching element in the right-hand side context is [merged][CopyableThreadContextElement.mergeForChild]
- * - Every CTCE from the right-hand side context that hasn't been merged is copied
- * - Everything else is added to the resulting context as is.
+
+ * The rules are as follows:
+ * - If both contexts have the same (by key) CTCE, they are [merged][CopyableThreadContextElement.mergeForChild].
+ * - If [isNewCoroutine] is `true`, the CTCEs that one context has and the other does not are
+ *   [copied][CopyableThreadContextElement.copyForChild].
+ * - If [isNewCoroutine] is `false`, then the CTCEs that the right context has and the left does not are copied,
+ *   but those that only the left context has are not copied but added to the resulting context as is.
+ * - Every non-CTCE is added to the resulting context as is.
  */
 private fun foldCopies(originalContext: CoroutineContext, appendContext: CoroutineContext, isNewCoroutine: Boolean): CoroutineContext {
     // Do we have something to copy left-hand side?
diff --git a/kotlinx-coroutines-core/jvm/test/CoroutineScopeTestJvm.kt b/kotlinx-coroutines-core/jvm/test/CoroutineScopeTestJvm.kt
@@ -0,0 +1,74 @@
+package kotlinx.coroutines
+
+import kotlinx.coroutines.testing.*
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class CoroutineScopeTestJvm: TestBase() {
+    /**
+     * Test the documented behavior of [CoroutineScope.newCoroutineContext] regarding the copyable context elements.
+     */
+    @Test
+    fun testNewCoroutineContextCopyableContextElements() {
+        val ce1L = MyMutableContextElement("key1", "value1_l")
+        val ce2L = MyMutableContextElement("key2", "value2_l")
+        val ce2R = MyMutableContextElement("key2", "value2_r")
+        val ce3R = MyMutableContextElement("key3", "value3_r")
+        val nonce1L = CoroutineExceptionHandler { _, _ -> }
+        val nonce2L = Dispatchers.Default
+        val nonce2R = Dispatchers.IO
+        val nonce3R = CoroutineName("name3_r")
+        val leftContext = randomlyShuffledContext(ce1L, ce2L, nonce1L, nonce2L)
+        val rightContext = randomlyShuffledContext(ce2R, ce3R, nonce2R, nonce3R)
+        CoroutineScope(leftContext).newCoroutineContext(rightContext).let { ctx ->
+            assertEquals("Copy of 'value1_l'", ctx[MyMutableContextElementKey("key1")]?.value)
+            assertEquals("Merged 'value2_l' and 'value2_r'", ctx[MyMutableContextElementKey("key2")]?.value)
+            assertEquals("Copy of 'value3_r'", ctx[MyMutableContextElementKey("key3")]?.value)
+            assertSame(nonce1L, ctx[CoroutineExceptionHandler])
+            assertSame(nonce2R, ctx[ContinuationInterceptor])
+            assertSame(nonce3R, ctx[CoroutineName])
+        }
+    }
+
+    private fun randomlyShuffledContext(
+        vararg elements: CoroutineContext.Element
+    ): CoroutineContext = elements.toList().shuffled().fold(EmptyCoroutineContext, CoroutineContext::plus)
+}
+
+class MyMutableContextElementKey(val key: String): CoroutineContext.Key<MyMutableContextElement> {
+    override fun equals(other: Any?): Boolean =
+        this === other || other is MyMutableContextElementKey && key == other.key
+
+    override fun hashCode(): Int = key.hashCode()
+}
+
+class MyMutableContextElement(
+    val keyId: String,
+    var value: String
+) : AbstractCoroutineContextElement(MyMutableContextElementKey(keyId)), CopyableThreadContextElement<String> {
+    override fun updateThreadContext(context: CoroutineContext): String {
+        return value
+    }
+
+    override fun restoreThreadContext(context: CoroutineContext, oldState: String) {
+        value = oldState
+    }
+
+    override fun toString(): String {
+        return "MyMutableContextElement(keyId='$keyId', value='$value')"
+    }
+
+    override fun equals(other: Any?): Boolean =
+        this === other || other is MyMutableContextElement && keyId == other.keyId && value == other.value
+
+    override fun hashCode(): Int = 31 * key.hashCode() + value.hashCode()
+
+    override fun copyForChild(): CopyableThreadContextElement<String> =
+        MyMutableContextElement(keyId, "Copy of '$value'")
+
+    override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext =
+        MyMutableContextElement(
+            keyId,
+            "Merged '$value' and '${(overwritingElement as MyMutableContextElement).value}'"
+        )
+}
diff --git a/kotlinx-coroutines-core/jvm/test/ThreadContextElementTest.kt b/kotlinx-coroutines-core/jvm/test/ThreadContextElementTest.kt
@@ -237,17 +237,53 @@ class ThreadContextElementTest : TestBase() {
 
     @Test
     fun testThreadLocalFlowOn() = runTest {
-        val myData = MyData()
-        myThreadLocal.set(myData)
-        expect(1)
-        flow {
-            assertEquals(myData, myThreadLocal.get())
-            emit(1)
+        val parameters: List<Triple<CoroutineContext, Boolean, Boolean>> =
+            listOf(EmptyCoroutineContext, Dispatchers.Default, Dispatchers.Unconfined).flatMap { dispatcher ->
+                listOf(true, false).flatMap { doYield ->
+                    listOf(true, false).map { useThreadLocalInOuterContext ->
+                        Triple(dispatcher, doYield, useThreadLocalInOuterContext)
+                    }
+                }
+            }
+        for ((dispatcher, doYield, useThreadLocalInOuterContext) in parameters) {
+            try {
+                testThreadLocalFlowOn(dispatcher, doYield, useThreadLocalInOuterContext)
+            } catch (e: Throwable) {
+                throw AssertionError("Failed with parameters: dispatcher=$dispatcher, " +
+                    "doYield=$doYield, " +
+                    "useThreadLocalInOuterContext=$useThreadLocalInOuterContext", e)
+            }
+        }
+    }
+
+    private fun testThreadLocalFlowOn(
+        extraFlowOnContext: CoroutineContext, doYield: Boolean, useThreadLocalInOuterContext: Boolean
+    ) = runTest {
+        try {
+            val myData1 = MyData()
+            val myData2 = MyData()
+            myThreadLocal.set(myData1)
+            withContext(if (useThreadLocalInOuterContext) myThreadLocal.asContextElement() else EmptyCoroutineContext) {
+                assertEquals(myData1, myThreadLocal.get())
+                flow {
+                    repeat(5) {
+                        assertEquals(myData2, myThreadLocal.get())
+                        emit(1)
+                        if (doYield) yield()
+                    }
+                }
+                    .flowOn(myThreadLocal.asContextElement(myData2) + extraFlowOnContext)
+                    .collect {
+                        if (useThreadLocalInOuterContext) {
+                            assertEquals(myData1, myThreadLocal.get())
+                        }
+                    }
+                assertEquals(myData1, myThreadLocal.get())
+            }
+            assertEquals(myData1, myThreadLocal.get())
+        } finally {
+            myThreadLocal.set(null)
         }
-            .flowOn(myThreadLocal.asContextElement() + Dispatchers.Default)
-            .single()
-        myThreadLocal.set(null)
-        finish(2)
     }
 }
 
diff --git a/kotlinx-coroutines-core/jvm/test/ThreadContextMutableCopiesTest.kt b/kotlinx-coroutines-core/jvm/test/ThreadContextMutableCopiesTest.kt
@@ -132,29 +132,27 @@ class ThreadContextMutableCopiesTest : TestBase() {
 
     @Test
     fun testDataIsCopiedThroughFlowOnUndispatched() = runTest {
-        expect(1)
-        val root = MyMutableElement(ArrayList())
-        val originalData = root.mutableData
+        val originalData = mutableListOf("X")
+        val root = MyMutableElement(originalData)
         flow {
             assertNotSame(originalData, threadLocalData.get())
+            assertEquals(originalData, threadLocalData.get())
             emit(1)
         }
             .flowOn(root)
             .single()
-        finish(2)
     }
 
     @Test
     fun testDataIsCopiedThroughFlowOnDispatched() = runTest {
-        expect(1)
-        val root = MyMutableElement(ArrayList())
-        val originalData = root.mutableData
+        val originalData = mutableListOf("X")
+        val root = MyMutableElement(originalData)
         flow {
             assertNotSame(originalData, threadLocalData.get())
+            assertEquals(originalData, threadLocalData.get())
             emit(1)
         }
             .flowOn(root + Dispatchers.Default)
             .single()
-        finish(2)
     }
 }
diff --git a/reactive/kotlinx-coroutines-reactive/src/ReactiveFlow.kt b/reactive/kotlinx-coroutines-reactive/src/ReactiveFlow.kt
@@ -69,19 +69,14 @@ private class PublisherAsFlow<T : Any>(
 
     override suspend fun collect(collector: FlowCollector<T>) {
         val collectContext = coroutineContext
-        val newDispatcher = context[ContinuationInterceptor]
-        if (newDispatcher == null || newDispatcher == collectContext[ContinuationInterceptor]) {
-            // fast path -- subscribe directly in this dispatcher
-            return collectImpl(collectContext + context, collector)
+        val newContext = collectContext.newCoroutineContext(context)
+        // quickest path: if the context has not changed, just subscribe inline
+        if (newContext == collectContext) {
+            return collectImpl(collectContext, collector)
         }
+        // TODO: copy-paste/share the ChannelFlowOperatorImpl implementation for the same-dispatcher quick path
         // slow path -- produce in a separate dispatcher
-        collectSlowPath(collector)
-    }
-
-    private suspend fun collectSlowPath(collector: FlowCollector<T>) {
-        coroutineScope {
-            collector.emitAll(produceImpl(this + context))
-        }
+        super.collect(collector)
     }
 
     private suspend fun collectImpl(injectContext: CoroutineContext, collector: FlowCollector<T>) {
diff --git a/reactive/kotlinx-coroutines-reactive/test/PublisherAsFlowThreadContextElementTest.kt b/reactive/kotlinx-coroutines-reactive/test/PublisherAsFlowThreadContextElementTest.kt
