Propagation of the coroutine context of await calls into Mono/Flux builders

Fixes #284

Author: SokolovaMaria

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

@@ -14,6 +14,10 @@ public final class kotlinx/coroutines/reactive/ChannelKt {
 	public static synthetic fun openSubscription$default (Lorg/reactivestreams/Publisher;IILjava/lang/Object;)Lkotlinx/coroutines/channels/ReceiveChannel;
 }
 
+public abstract interface class kotlinx/coroutines/reactive/ContextInjector {
+	public abstract fun injectCoroutineContext (Lorg/reactivestreams/Publisher;Lkotlin/coroutines/CoroutineContext;)Lorg/reactivestreams/Publisher;
+}
+
 public final class kotlinx/coroutines/reactive/ConvertKt {
 	public static final fun asPublisher (Lkotlinx/coroutines/channels/ReceiveChannel;Lkotlin/coroutines/CoroutineContext;)Lorg/reactivestreams/Publisher;
 	public static synthetic fun asPublisher$default (Lkotlinx/coroutines/channels/ReceiveChannel;Lkotlin/coroutines/CoroutineContext;ILjava/lang/Object;)Lorg/reactivestreams/Publisher;

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

@@ -28,6 +28,11 @@ 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/Await.kt

@@ -10,6 +10,7 @@ import kotlinx.coroutines.suspendCancellableCoroutine
 import org.reactivestreams.Publisher
 import org.reactivestreams.Subscriber
 import org.reactivestreams.Subscription
+import java.util.*
 import kotlin.coroutines.*
 
 /**
@@ -22,7 +23,7 @@ import kotlin.coroutines.*
  *
  * @throws NoSuchElementException if publisher does not emit any value
  */
-public suspend fun <T> Publisher<T>.awaitFirst(): T = awaitOne(Mode.FIRST)
+public suspend fun <T> Publisher<T>.awaitFirst(): T = injectCoroutineContext().awaitOne(Mode.FIRST)
 
 /**
  * Awaits for the first value from the given observable or the [default] value if none is emitted without blocking a
@@ -32,7 +33,8 @@ public suspend fun <T> Publisher<T>.awaitFirst(): T = awaitOne(Mode.FIRST)
  * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException].
  */
-public suspend fun <T> Publisher<T>.awaitFirstOrDefault(default: T): T = awaitOne(Mode.FIRST_OR_DEFAULT, default)
+public suspend fun <T> Publisher<T>.awaitFirstOrDefault(default: T): T =
+    injectCoroutineContext().awaitOne(Mode.FIRST_OR_DEFAULT, default)
 
 /**
  * Awaits for the first value from the given observable or `null` value if none is emitted without blocking a
@@ -42,7 +44,7 @@ public suspend fun <T> Publisher<T>.awaitFirstOrDefault(default: T): T = awaitOn
  * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException].
  */
-public suspend fun <T> Publisher<T>.awaitFirstOrNull(): T? = awaitOne(Mode.FIRST_OR_DEFAULT)
+public suspend fun <T> Publisher<T>.awaitFirstOrNull(): T? = injectCoroutineContext().awaitOne(Mode.FIRST_OR_DEFAULT)
 
 /**
  * Awaits for the first value from the given observable or call [defaultValue] to get a value if none is emitted without blocking a
@@ -52,7 +54,8 @@ public suspend fun <T> Publisher<T>.awaitFirstOrNull(): T? = awaitOne(Mode.FIRST
  * If the [Job] of the current coroutine is cancelled or completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException].
  */
-public suspend fun <T> Publisher<T>.awaitFirstOrElse(defaultValue: () -> T): T = awaitOne(Mode.FIRST_OR_DEFAULT) ?: defaultValue()
+public suspend fun <T> Publisher<T>.awaitFirstOrElse(defaultValue: () -> T): T =
+    injectCoroutineContext().awaitOne(Mode.FIRST_OR_DEFAULT) ?: defaultValue()
 
 /**
  * Awaits for the last value from the given publisher without blocking a thread and
@@ -64,7 +67,7 @@ public suspend fun <T> Publisher<T>.awaitFirstOrElse(defaultValue: () -> T): T =
  *
  * @throws NoSuchElementException if publisher does not emit any value
  */
-public suspend fun <T> Publisher<T>.awaitLast(): T = awaitOne(Mode.LAST)
+public suspend fun <T> Publisher<T>.awaitLast(): T = injectCoroutineContext().awaitOne(Mode.LAST)
 
 /**
  * Awaits for the single value from the given publisher without blocking a thread and
@@ -77,10 +80,21 @@ public suspend fun <T> Publisher<T>.awaitLast(): T = awaitOne(Mode.LAST)
  * @throws NoSuchElementException if publisher does not emit any value
  * @throws IllegalArgumentException if publisher emits more than one value
  */
-public suspend fun <T> Publisher<T>.awaitSingle(): T = awaitOne(Mode.SINGLE)
+public suspend fun <T> Publisher<T>.awaitSingle(): T = injectCoroutineContext().awaitOne(Mode.SINGLE)
 
 // ------------------------ private ------------------------
 
+// ContextInjector service is implemented in `kotlinx-coroutines-reactor` module only.
+// If `kotlinx-coroutines-reactor` module is not included, the list is empty.
+private val contextInjectors: List<ContextInjector> =
+    ServiceLoader.load(ContextInjector::class.java, ContextInjector::class.java.classLoader).toList()
+
+private suspend fun <T> Publisher<T>.injectCoroutineContext(): Publisher<T> {
+    return if (contextInjectors.isNotEmpty()) {
+        contextInjectors[0].injectCoroutineContext(this, coroutineContext)
+    } else this
+}
+
 private enum class Mode(val s: String) {
     FIRST("awaitFirst"),
     FIRST_OR_DEFAULT("awaitFirstOrDefault"),

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

@@ -0,0 +1,14 @@
+package kotlinx.coroutines.reactive
+
+import kotlinx.coroutines.InternalCoroutinesApi
+import org.reactivestreams.Publisher
+import kotlin.coroutines.CoroutineContext
+
+/** @suppress */
+@InternalCoroutinesApi
+public interface ContextInjector {
+    /**
+     * Injects the coroutine context into the context of the publisher.
+     */
+    public fun <T> injectCoroutineContext(publisher: Publisher<T>, coroutineContext: CoroutineContext): Publisher<T>
+}

reactive/kotlinx-coroutines-reactor/resources/META-INF/services/kotlinx.coroutines.reactive.ContextInjector

@@ -0,0 +1 @@
+kotlinx.coroutines.reactor.ReactorContextInjector

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

@@ -30,6 +30,22 @@ import kotlin.coroutines.*
  *    .subscribe()
  * }
  * ```
+ *
+ * [CoroutineContext] of a suspendable function that awaits a value from [Mono] or [Flux] instance
+ * is propagated into [mono] and [flux] Reactor builders.
+ *
+ * Example usage:
+ *
+ * ```
+ * launch(Context.of("key", "value").asCoroutineContext()) {
+ *   assertEquals(bar().awaitFirst(), "value")
+ * }
+ *
+ * fun bar(): Mono<String> = mono {
+ *   coroutineContext[ReactorContext]!!.context.get("key")
+ * }
+ * ```
+}
  */
 @ExperimentalCoroutinesApi
 public class ReactorContext(val context: Context) : AbstractCoroutineContextElement(ReactorContext) {

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

@@ -0,0 +1,28 @@
+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
+
+/** @suppress */
+@InternalCoroutinesApi
+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].
+     */
+    override fun <T> injectCoroutineContext(publisher: Publisher<T>, coroutineContext: CoroutineContext): Publisher<T> {
+        val reactorContext = coroutineContext[ReactorContext]?.context
+        return when(publisher) {
+            is Mono ->
+                publisher.subscriberContext(reactorContext ?: Context.empty())
+            is Flux ->
+                publisher.subscriberContext(reactorContext ?: Context.empty())
+            else -> publisher
+        }
+    }
+}

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

@@ -3,6 +3,7 @@ package kotlinx.coroutines.reactor
 import kotlinx.coroutines.*
 import kotlinx.coroutines.reactive.*
 import org.junit.Test
+import reactor.core.publisher.*
 import reactor.util.context.Context
 import kotlin.test.assertEquals
 
@@ -14,8 +15,8 @@ class ReactorContextTest {
             buildString {
                 (1..7).forEach { append(ctx?.getOrDefault(it, "noValue")) }
             }
-        }   .subscriberContext(Context.of(2, "2", 3, "3", 4, "4", 5, "5"))
-            .subscriberContext { ctx -> ctx.put(6, "6") }
+        }  .subscriberContext(Context.of(2, "2", 3, "3", 4, "4", 5, "5"))
+           .subscriberContext { ctx -> ctx.put(6, "6") }
         assertEquals(mono.awaitFirst(), "1234567")
     }
 
@@ -29,4 +30,48 @@ class ReactorContextTest {
         var i = 0
         flux.subscribe { str -> i++; assertEquals(str, i.toString()) }
     }
+
+    @Test
+    fun testAwait() = runBlocking(Context.of(3, "3").asCoroutineContext()) {
+        val result = mono(Context.of(1, "1").asCoroutineContext()) {
+            val ctx = coroutineContext[ReactorContext]?.context
+            buildString {
+                (1..3).forEach { append(ctx?.getOrDefault(it, "noValue")) }
+            }
+        }  .subscriberContext(Context.of(2, "2"))
+            .awaitFirst()
+        assertEquals(result, "123")
+    }
+
+    @Test
+    fun testMonoAwaitContextPropagation() = runBlocking(Context.of(7, "7").asCoroutineContext()) {
+        assertEquals(m().awaitFirst(), "7")
+        assertEquals(m().awaitFirstOrDefault("noValue"), "7")
+        assertEquals(m().awaitFirstOrNull(), "7")
+        assertEquals(m().awaitFirstOrElse { "noValue" }, "7")
+        assertEquals(m().awaitLast(), "7")
+        assertEquals(m().awaitSingle(), "7")
+    }
+
+    @Test
+    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")
+        assertEquals(f().awaitFirstOrElse { "noValue" }, "1")
+        assertEquals(f().awaitLast(), "3")
+        var i = 0
+        f().subscribe { str -> i++; assertEquals(str, i.toString()) }
+    }
+
+    private fun m(): Mono<String> = mono {
+        val ctx = coroutineContext[ReactorContext]?.context
+        ctx?.getOrDefault(7, "noValue")
+    }
+
+
+    private fun f(): Flux<String?> = flux {
+        val ctx = coroutineContext[ReactorContext]?.context
+        (1..3).forEach { send(ctx?.getOrDefault(it, "noValue")) }
+    }
 }