ReactorContextInjector does not have contextCapture

[olga-larina]

Hello!
There are similar issues in other projects made by another author:
reactor/reactor-kotlin-extensions#59
reactor/reactor-core#3406

Maybe these changes should be made in https://github.com/Kotlin/kotlinx.coroutines/blob/f4f519b36734238ec686dfaec1e174086691781e/reactive/kotlinx-coroutines-reactor/src/ReactorContextInjector.kt
Please, take a look.

Use case

For users using kotlin coroutine, ThreadContextElement is used to propagate MDC, etc.

The block operator is supported by reactor/reactor-core#3420, but not the await function. Therefore, we are forced to write it this way.

// e.g. WebClient
webClient.post()
    .uri("...")
    // omitted
    .contextCapture()
    .awaitSingle()

It would be great if the await function is also supported, like the block operator.

The Shape of the API

This option can be achieved by adding contextCapture() in ContextInjector. It looks like we could add it to existing ReactorContextInjector instead of creating new ContextInjector.
Possible solution may look like:

package kotlinx.coroutines.reactor

import kotlinx.coroutines.reactive.*
import org.reactivestreams.*
import reactor.core.publisher.*
import reactor.util.context.*
import kotlin.coroutines.*

internal class ReactorContextInjector : ContextInjector {
    
    override fun <T> injectCoroutineContext(publisher: Publisher<T>, coroutineContext: CoroutineContext): Publisher<T> {
        val reactorContext = coroutineContext[ReactorContext]?.context
        return when(publisher) {
            is Mono -> publisher.contextWriteNullable(reactorContext).contextCapture()
            is Flux -> publisher.contextWriteNullable(reactorContext).contextCapture()
            else -> publisher
        }
    }

    private fun <T> Mono<T>.contextWriteNullable(reactorContext: Context?) = if (reactorContext == null) {
        this
    } else {
        this.contextWrite(reactorContext)
    }

    private fun <T> Flux<T>.contextWriteNullable(reactorContext: Context?) = if (reactorContext == null) {
        this
    } else {
        this.contextWrite(reactorContext)
    }
}

[dkhalanskyjb]

Hello!

1. Let's double-check first what behavior this is supposed to achieve: the way I understand it, you would like the context of the caller of await to influence the context of the caller publisher. Is that right?
2. If yes, then another question: is it true that calling .asFlow().singleOrNull() on Mono instead of .await() already has the desired behavior?

[olga-larina]

1. I would like to save the context of the caller between threads while using Reactor and awaitSingle.
   Now (without contextCapture()) MDC and values of the thread locals are lost (with block() all values are saved without contextCapture()).

                webClient.get().uri("https://example.com").retrieve().toBodilessEntity()
                    .doOnSuccess {
                        log.info(
                            "doOnSuccess. MDC={}, myThreadLocal={}",
                            MDC.getCopyOfContextMap(), myThreadLocal.get()
                        )
                    }
                    // .contextCapture()
                    .awaitSingle()

2. No. Using .asFlow().singleOrNull() also looses context. Only block() saves context. Reproducer was made by another person: https://github.com/be-hase/context-propagation-report-202309/blob/main/coroutines-issue/src/main/kotlin/example/CoroutinesIssueApplication.kt

[olga-larina]

Also the dependency io.projectreactor:reactor-core:3.4.1 should be upgraded to 3.6.0 (or higher). There are the issues on reactor which show that there are the problems with the context propagation on the old versions. And contextCapture() has been introduced later.
Maybe it will be enough to add only contextCapture() instead of contextWrite(reactorContext).contextCapture() because contextCapture() should do all work automatically.
But using contextWrite(reactorContext).contextCapture() looks safer.