fixup! Reactor coroutine context propagation in more places

Author: elizarov

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-reactor.txt

@@ -32,11 +32,6 @@ public final class kotlinx/coroutines/reactor/ReactorContext : kotlin/coroutines
 public final class kotlinx/coroutines/reactor/ReactorContext$Key : kotlin/coroutines/CoroutineContext$Key {
 }
 
-public final class kotlinx/coroutines/reactor/ReactorContextInjector : kotlinx/coroutines/reactive/ContextInjector {
-	public fun <init> ()V
-	public fun injectCoroutineContext (Lorg/reactivestreams/Publisher;Lkotlin/coroutines/CoroutineContext;)Lorg/reactivestreams/Publisher;
-}
-
 public final class kotlinx/coroutines/reactor/ReactorContextKt {
 	public static final fun asCoroutineContext (Lreactor/util/context/Context;)Lkotlinx/coroutines/reactor/ReactorContext;
 }

reactive/kotlinx-coroutines-reactive/src/PublisherAsFlow.kt

@@ -47,7 +47,9 @@ private class PublisherAsFlow<T : Any>(
         // use another channel for conflation (cannot do openSubscription)
         if (capacity < 0) return super.produceImpl(scope)
         // Open subscription channel directly
-        val channel = publisher.openSubscription(capacity)
+        val channel = publisher
+            .injectCoroutineContext(scope.coroutineContext)
+            .openSubscription(capacity)
         val handle = scope.coroutineContext[Job]?.invokeOnCompletion(onCancelling = true) { cause ->
             channel.cancel(cause?.let {
                 it as? CancellationException ?: CancellationException("Job was cancelled", it)

reactive/kotlinx-coroutines-reactor/src/ReactorContextInjector.kt

@@ -1,16 +1,12 @@
 package kotlinx.coroutines.reactor
 
-import kotlinx.coroutines.InternalCoroutinesApi
-import kotlinx.coroutines.reactive.ContextInjector
-import org.reactivestreams.Publisher
-import reactor.core.publisher.Flux
-import reactor.core.publisher.Mono
-import reactor.util.context.Context
-import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.reactive.*
+import org.reactivestreams.*
+import reactor.core.publisher.*
+import reactor.util.context.*
+import kotlin.coroutines.*
 
-/** @suppress */
-@InternalCoroutinesApi
-class ReactorContextInjector : ContextInjector {
+internal class ReactorContextInjector : ContextInjector {
     /**
      * Injects all values from the [ReactorContext] entry of the given coroutine context
      * into the downstream [Context] of Reactor's [Publisher] instances of [Mono] or [Flux].

reactive/kotlinx-coroutines-reactor/test/ReactorContextTest.kt

@@ -1,12 +1,13 @@
 package kotlinx.coroutines.reactor
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.reactive.*
 import org.junit.Test
 import reactor.core.publisher.*
-import reactor.util.context.Context
-import kotlin.test.assertEquals
-import kotlinx.coroutines.flow.*
+import reactor.util.context.*
+import kotlin.test.*
 
 class ReactorContextTest {
     @Test
@@ -55,7 +56,9 @@ class ReactorContextTest {
     }
 
     @Test
-    fun testFluxAwaitContextPropagation() = runBlocking<Unit>(Context.of(1, "1", 2, "2", 3, "3").asCoroutineContext()) {
+    fun testFluxAwaitContextPropagation() = runBlocking<Unit>(
+        Context.of(1, "1", 2, "2", 3, "3").asCoroutineContext()
+    ) {
         assertEquals(f().awaitFirst(), "1")
         assertEquals(f().awaitFirstOrDefault("noValue"), "1")
         assertEquals(f().awaitFirstOrNull(), "1")
@@ -77,18 +80,32 @@ class ReactorContextTest {
     }
 
     @Test
-    fun testFlowToFluxContextPropagation() = runBlocking(Context.of(1, "1", 2, "2", 3, "3").asCoroutineContext()) {
+    fun testFlowToFluxContextPropagation() = runBlocking(
+        Context.of(1, "1", 2, "2", 3, "3").asCoroutineContext()
+    ) {
         var i = 0
+        // call "collect" on the converted Flow
         bar().collect { str ->
             i++; assertEquals(str, i.toString())
         }
         assertEquals(i, 3)
     }
 
-    suspend fun bar(): Flow<String> {
-        return flux {
-            val ctx = coroutineContext[ReactorContext]!!.context
-            (1..3).forEach { send(ctx.getOrDefault(it, "noValue")) }
-        }.asFlow()
+    @Test
+    fun testFlowToFluxDirectContextPropagation() = runBlocking(
+        Context.of(1, "1", 2, "2", 3, "3").asCoroutineContext()
+    ) {
+        var i = 0
+        // convert resulting flow to channel using "produceIn"
+        val channel = bar().produceIn(this)
+        channel.consumeEach { str ->
+            i++; assertEquals(str, i.toString())
+        }
+        assertEquals(i, 3)
     }
+
+    private fun bar(): Flow<String> = flux {
+        val ctx = coroutineContext[ReactorContext]!!.context
+        (1..3).forEach { send(ctx.getOrDefault(it, "noValue")) }
+    }.asFlow()
 }