Fix thread context elements not being cleaned up after collect

Wait for the upstream coroutine to finish before restoring the
thread context element state to the old one.

Fixes #4403

Author: dkhalanskyjb

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>(

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)
     }
 }