Context passing between coroutines and Reactor Mono/Flux

Fixes #284

Author: SokolovaMaria

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

@@ -23,3 +23,20 @@ public final class kotlinx/coroutines/reactive/PublishKt {
 	public static synthetic fun publish$default (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/reactivestreams/Publisher;
 }
 
+public final class kotlinx/coroutines/reactive/PublisherCoroutine : kotlinx/coroutines/AbstractCoroutine, kotlinx/coroutines/channels/ProducerScope, kotlinx/coroutines/selects/SelectClause2, org/reactivestreams/Subscription {
+	public fun <init> (Lkotlin/coroutines/CoroutineContext;Lorg/reactivestreams/Subscriber;)V
+	public fun cancel ()V
+	public fun close (Ljava/lang/Throwable;)Z
+	public fun getChannel ()Lkotlinx/coroutines/channels/SendChannel;
+	public fun getOnSend ()Lkotlinx/coroutines/selects/SelectClause2;
+	public fun invokeOnClose (Lkotlin/jvm/functions/Function1;)Ljava/lang/Void;
+	public synthetic fun invokeOnClose (Lkotlin/jvm/functions/Function1;)V
+	public fun isClosedForSend ()Z
+	public fun isFull ()Z
+	public fun offer (Ljava/lang/Object;)Z
+	public synthetic fun onCompleted (Ljava/lang/Object;)V
+	public fun registerSelectClause2 (Lkotlinx/coroutines/selects/SelectInstance;Ljava/lang/Object;Lkotlin/jvm/functions/Function2;)V
+	public fun request (J)V
+	public fun send (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+

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

@@ -15,6 +15,19 @@ public final class kotlinx/coroutines/reactor/MonoKt {
 	public static synthetic fun mono$default (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lreactor/core/publisher/Mono;
 }
 
+public final class kotlinx/coroutines/reactor/ReactorContext : kotlin/coroutines/AbstractCoroutineContextElement {
+	public static final field Key Lkotlinx/coroutines/reactor/ReactorContext$Key;
+	public fun <init> (Lreactor/util/context/Context;)V
+	public final fun getContext ()Lreactor/util/context/Context;
+}
+
+public final class kotlinx/coroutines/reactor/ReactorContext$Key : kotlin/coroutines/CoroutineContext$Key {
+}
+
+public final class kotlinx/coroutines/reactor/ReactorContextKt {
+	public static final fun asCoroutineContext (Lreactor/util/context/Context;)Lkotlin/coroutines/CoroutineContext;
+}
+
 public final class kotlinx/coroutines/reactor/SchedulerCoroutineDispatcher : kotlinx/coroutines/CoroutineDispatcher, kotlinx/coroutines/Delay {
 	public fun <init> (Lreactor/core/scheduler/Scheduler;)V
 	public fun delay (JLkotlin/coroutines/Continuation;)Ljava/lang/Object;

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

@@ -54,7 +54,8 @@ private const val CLOSED = -1L    // closed, but have not signalled onCompleted/
 private const val SIGNALLED = -2L  // already signalled subscriber onCompleted/onError
 
 @Suppress("CONFLICTING_JVM_DECLARATIONS", "RETURN_TYPE_MISMATCH_ON_INHERITANCE")
-private class PublisherCoroutine<in T>(
+@InternalCoroutinesApi
+public class PublisherCoroutine<in T>(
     parentContext: CoroutineContext,
     private val subscriber: Subscriber<T>
 ) : AbstractCoroutine<Unit>(parentContext, true), ProducerScope<T>, Subscription, SelectClause2<T, SendChannel<T>> {

reactive/kotlinx-coroutines-reactor/build.gradle

@@ -11,4 +11,12 @@ tasks.withType(dokka.getClass()) {
     externalDocumentationLink {
         url = new URL('https://projectreactor.io/docs/core/3.2.5.RELEASE/api/')
     }
+}
+
+compileTestKotlin {
+    kotlinOptions.jvmTarget = "1.8"
+}
+
+compileKotlin {
+    kotlinOptions.jvmTarget = "1.8"
 }

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

@@ -4,11 +4,13 @@
 
 package kotlinx.coroutines.reactor
 
+import kotlin.coroutines.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.reactive.*
+import org.reactivestreams.*
 import reactor.core.publisher.*
-import kotlin.coroutines.*
+import reactor.core.*
 
 /**
  * Creates cold reactive [Flux] that runs a given [block] in a coroutine.
@@ -32,9 +34,25 @@ import kotlin.coroutines.*
  * **Note: This is an experimental api.** Behaviour of publishers that work as children in a parent scope with respect
  *        to cancellation and error handling may change in the future.
  */
+
 @ExperimentalCoroutinesApi
 fun <T> CoroutineScope.flux(
     context: CoroutineContext = EmptyCoroutineContext,
     @BuilderInference block: suspend ProducerScope<T>.() -> Unit
-): Flux<T> =
-    Flux.from(publish(newCoroutineContext(context), block = block))
+): Flux<T> = Flux.from(reactorPublish(newCoroutineContext(context), block = block))
+
+
+private fun <T> CoroutineScope.reactorPublish(
+    context: CoroutineContext = EmptyCoroutineContext,
+    @BuilderInference block: suspend ProducerScope<T>.() -> Unit
+): Publisher<T> = Publisher { subscriber ->
+    // specification requires NPE on null subscriber
+    if (subscriber == null) throw NullPointerException("Subscriber cannot be null")
+    require(subscriber is CoreSubscriber) { "Subscriber is not an instance of CoreSubscriber, context can not be extracted." }
+    val currentContext = subscriber.currentContext()
+    val reactorContext = (coroutineContext[ReactorContext]?.context?.putAll(currentContext) ?: currentContext).asCoroutineContext()
+    val newContext = newCoroutineContext(context + reactorContext)
+    val coroutine = PublisherCoroutine(newContext, subscriber)
+    subscriber.onSubscribe(coroutine) // do it first (before starting coroutine), to avoid unnecessary suspensions
+    coroutine.start(CoroutineStart.DEFAULT, coroutine, block)
+}

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

@@ -32,7 +32,8 @@ fun <T> CoroutineScope.mono(
     context: CoroutineContext = EmptyCoroutineContext,
     block: suspend CoroutineScope.() -> T?
 ): Mono<T> = Mono.create { sink ->
-    val newContext = newCoroutineContext(context)
+    val reactorContext = (coroutineContext[ReactorContext]?.context?.putAll(sink.currentContext()) ?: sink.currentContext()).asCoroutineContext()
+    val newContext = newCoroutineContext(context + reactorContext)
     val coroutine = MonoCoroutine(newContext, sink)
     sink.onDispose(coroutine)
     coroutine.start(CoroutineStart.DEFAULT, coroutine, block)
@@ -53,12 +54,11 @@ private class MonoCoroutine<in T>(
     override fun onCompletedExceptionally(exception: Throwable) {
         if (!disposed) sink.error(exception)
     }
-    
+
     override fun dispose() {
         disposed = true
         cancel()
     }
 
     override fun isDisposed(): Boolean = disposed
 }
-

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

@@ -0,0 +1,50 @@
+package kotlinx.coroutines.reactor
+
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import reactor.util.context.Context
+import kotlin.coroutines.*
+
+/**
+ * Marks coroutine context element that contains Reactor's [Context] elements in [context] for seamless integration
+ * between [CoroutineContext] and Reactor's [Context].
+ *
+ * [Context.asCoroutineContext] is defined to add Reactor's [Context] elements as part of [CoroutineContext].
+ *
+ * Reactor builders: [mono], [flux] can extract the reactor context from their coroutine context and
+ * pass it on. Modifications of reactor context can be retrieved by `coroutineContext[ReactorContext]`.
+ *
+ * Example usage:
+ *
+ * Passing reactor context from coroutine builder to reactor entity:
+ *
+ * ```
+ * launch(Context.of("key", "value").asCoroutineContext()) {
+ *   mono {
+ *     assertEquals(coroutineContext[ReactorContext]!!.context.get("key"), "value")
+ *   }.subscribe()
+ * }
+ * ```
+ *
+ * Accessing modified reactor context enriched from downstream via coroutine context:
+ *
+ * ```
+ * launch {
+ *   mono {
+ *     assertEquals(coroutineContext[ReactorContext]!!.context.get("key"), "value")
+ *   }.subscriberContext(Context.of("key", "value"))
+ *    .subscribe()
+ * }
+ * ```
+ */
+@ExperimentalCoroutinesApi
+public class ReactorContext(val context: Context) : AbstractCoroutineContextElement(ReactorContext) {
+    companion object Key : CoroutineContext.Key<ReactorContext>
+}
+
+
+/**
+ * Wraps the given [Context] into [ReactorContext], so it can be added to coroutine's context
+ * and later retrieved via `coroutineContext[ReactorContext]`.
+ */
+@ExperimentalCoroutinesApi
+public fun Context.asCoroutineContext(): CoroutineContext = ReactorContext(this)

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

@@ -0,0 +1,54 @@
+package kotlinx.coroutines.reactor
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.reactive.*
+import org.junit.Test
+import reactor.util.context.Context
+import kotlin.test.assertEquals
+
+class ReactorContextTest {
+    @Test
+    fun testMonoHookedContext() = runBlocking(Context.of(1, "1", 7, "7").asCoroutineContext()) {
+        val mono = mono {
+            val ctx = coroutineContext[ReactorContext]?.context
+            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") }
+        assertEquals(mono.awaitFirst(), "1234567")
+    }
+
+    @Test
+    fun testFluxContext() = runBlocking<Unit>(Context.of(1, "1", 7, "7").asCoroutineContext()) {
+        val flux = flux<String?> {
+            val ctx = coroutineContext[ReactorContext]!!.context
+            (1..7).forEach { send(ctx.getOrDefault(it, "noValue")) }
+        }   .subscriberContext(Context.of(2, "2", 3, "3", 4, "4", 5, "5"))
+            .subscriberContext { ctx -> ctx.put(6, "6") }
+        var i = 0
+        flux.subscribe { str -> i++; assertEquals(str, i.toString()) }
+    }
+
+    @Test
+    fun testContextFromCoroutineBuilderToReactor() = runBlocking<Unit> {
+        launch(Context.of(1, "1").asCoroutineContext()) {
+            val mono = mono {
+                assertEquals(coroutineContext[ReactorContext]!!.context[1], "1")
+            }
+            mono.subscribe()
+        }
+    }
+
+    @Test
+    fun testContextFromReactorToCoroutine() = runBlocking<Unit>(Context.of(1, "global 1").asCoroutineContext()) {
+        mono {
+            assertEquals(coroutineContext[ReactorContext]!!.context[2], "2")
+            assertEquals(coroutineContext[ReactorContext]!!.context[1], "global 1")
+        }.subscriberContext(Context.of(2, "2", 3, "3"))
+                .subscribe()
+        mono {
+            assertEquals(coroutineContext[ReactorContext]!!.context.getOrDefault(1, "no value"), "global 1")
+        }.subscribe()
+    }
+}