Spring Webflux: CoRouterFunctionDsl should expose the used CoroutineContext

[FrontierPsychiatrist]

Affects: 5.3.7

---

I am using the CoRouterFunctionDsl to build a Spring Webflux application with suspend handlers in Kotlin. Recently I wanted to add some SLF4J MDC values (e.g. an AWS X-Ray Trace id) which requires the usage of https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-core/kotlinx.coroutines/-thread-context-element/. MDC is ThreadLocal based which does not play well with coroutines.

Adding such a ThreadContextElement can be done in a generic way (so it doesn't have to be repeated in each handler) in the private method asHandlerFunction in the class CoRouterFunctionDsl:

mono(Dispatchers.Unconfined) {
    init(it)
}

but there is no way of doing that. For now I copied the class and changed the function myself. The data I need there is based on the subscriber context of the Mono (coming from a WebFilter) and the ServerRequest, so it would be neat to inject something like coroutineContextProvider: (req: ServerRequest) -> Mono<CoroutineContext> into the DSL:

private val contextProvider = (req: ServerRequest) ->
  Mono.deferContextual { ctx ->
    // Extract data from req and ctx
    // Return a coroutine context
  }

coRouter(contextProvider) {
  GET("/info", ...)
}

// In CoRouterFunctionDsl

private fun asHandlerFunction(init: suspend (ServerRequest) -> ServerResponse) = HandlerFunction { req ->
    contextProvider(req)
      .flatMap { coroutineContext ->
              mono(coroutineContext) {
                  init(req)
              }
      }
}

I hope that makes sense. If there is an easier solution that I oversaw I am of course also open to that.

[poutsma]

@sdeleuze Could you take a look at this one, please?

[sdeleuze]

@poutsma Do we provide a way to do the same thing on Reactive side via Reactor context? If not, what would that look like and would it b something we want to implement?

[FrontierPsychiatrist]

@sdeleuze I think in a pure Reactor implementation this is easier since the call chain "stays inside the framework". It can be easily done by adding a WebFilter that uses Mono.contextWrite. Adding a filter in coRouter does not have the same effect as there will be another "switch" back to Reactor after the filter and the coroutine context will be "forgotten".

[sdeleuze]

I am not sure how this should be implemented and exposed, but indeed getting access to the CoroutineContext seems a valid use case.