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.

[sdeleuze]

I am considering adding a fun context(provider: suspend (ServerRequest) -> CoroutineContext) method to the DSL in Spring Framework 6.1 to provide a custom CoroutineContext in various places of the DSL. It would allow to write code like:

coRouter {
  context { request ->
    delay(1) // Can invoke suspending functions
    CoroutineName(request.headers().firstHeader("Custom-Header")!!)
  }
  GET("/") {
    // ...
  }
}

[FrontierPsychiatrist]

That seems like pretty much exactly what would help me! E.g. something like this is what I do at the moment)

coRouter {
  context { request ->
    MDC.put("traceId", req.header("traceId"))
    MDC.put("userId", req.awaitPrincipal().name)
    // Not sure if NonCancellable will work like this, in my tests the mono() creator did not accept it
    kotlinx.coroutines.slf4j.MDCContext() + NonCancellable
  }
  // mappings
}

[sdeleuze]

Merged, please try 6.1.0-SNAPSHOT builds once created by the CI and provide a feedback if possible.

[FrontierPsychiatrist]

I just tested it in a small application and it worked fine!