Spring Propagate CoroutineContext in CoWebFilter

Motivation

There are cases where the coroutine dispatcher is changed when performing blocking processing.

@GetMapping("/sample")
suspend fun sample(): String {
    // do non-blocking call

    withContext(blockingDispatcher) {
        // do blocking call
    }

    // do non-blocking call

    ...
}

But now WebFlux uses Unconfined dispatcher. https://github.com/spring-projects/spring-framework/blob/5.3.x/spring-core/src/main/java/org/springframework/core/CoroutinesUtils.java#L74

So once we change the dispatcher, subsequent processing will also be executed by that dispatcher. For example, in the following example, it will be executed by blockingExecutor thread even after withContext.

@GetMapping("/test/hello1")
suspend fun hello1(): String {
    // reactor-http-nio-N thread
    log.info("thread={}", Thread.currentThread().name)

    withContext(blockingDispatcher) {
        // blockingExecutor-N thread
        log.info("thread={}", Thread.currentThread().name)

        // do blocking call
    }

    // blockingExecutor-N thread
    log.info("thread={}", Thread.currentThread().name)

    return "hello"
}

We can work around this issue by using the default dispatcher instead of the Unconfined dispatcher.

@GetMapping("/test/hello2")
suspend fun hello2(): String {
    return withContext(Dispatchers.Default) {
        // DefaultDispatcher-worker-N thread
        log.info("thread={}", Thread.currentThread().name)

        withContext(blockingDispatcher) {
            // blockingExecutor-N thread
            log.info("thread={}", Thread.currentThread().name)

            // do blocking call
        }

        // DefaultDispatcher-worker-N thread
        log.info("thread={}", Thread.currentThread().name)

        "hello"
    }
}

However, writing withContext(Dispatchers.Default) on all controller methods can be tedious. So I want to be able to change the coroutine dispatcher used by WebFlux.

How

How about making it possible to change by registering a class such as CoroutinesDispatchersProvider in the bean?

interface CoroutinesDispatchersProvider {
    fun provide(request: ServerHttpRequest): CoroutineDispatcher
}

For example, a framework called armeria allows we to specify a dispatcher https://armeria.dev/docs/server-annotated-service#coroutine-dispatcher

Comment From: be-hase

To give you another example, grpc-kotlin supports customizing coroutine context. https://github.com/grpc/grpc-kotlin/issues/66

Being able to customize the context makes it much easier to propagate MDCs and Brave's Tracing. e.g. https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-slf4j/kotlinx.coroutines.slf4j/-m-d-c-context/index.html

Please consider it.

Comment From: zsiegel

This would be very helpful in regards to MDC context as mentioned above.

I would be happy to put together a PR if any of the team members can point me in a direction that might be suitable to put this configuration?

Comment From: be-hase

This would be very helpful in regards to MDC context as mentioned above.

Yes. This is very ver helpful for kotlin coroutine user.

Comment From: 2hyjun

~This feature has been implemented on this commit~

Invoking suspend functions with customized CoroutineContext has been implemented, but there's no way to provide it currently.

Comment From: wplong11

@2hyjun This feature is not implemented. That commit only added new invokeSuspendingFunction that accepts CoroutineContext. To implement this feature, InvocableHandlerMethod must be fixed to use new invokeSuspendingFunction and CoroutinesDispatchersProvider should be added.

https://github.com/spring-projects/spring-framework/blob/983c6e233fa21e2672484ee8f4f8e8c8bcd51f95/spring-webflux/src/main/java/org/springframework/web/reactive/result/method/InvocableHandlerMethod.java#L139-L145

Comment From: wplong11

@rstoyanchev @be-hase Can we discuss about design of CoroutinesDispatchersProvider? If there is an agreement on the design, I am willing to write a PoC PR.

Comment From: be-hase

@wplong11 Sorry for the late reply. of course! always ready to discuss :)

Comment From: efemoney

It should be a HandlerMethodCoroutineContextProvider and not just a dispatcher provider. To implement things like CoroutineNameing endpoint calls, adding custom context elements and also pave the way for 3p to integrate

Comment From: sdeleuze

Let's try to see if we can leverage CoWebFilter to support this use case by taking inspiration of https://github.com/spring-projects/spring-framework/issues/26977#issuecomment-1696376264.

Comment From: be-hase

Thank you. But you must be using router function style, right...? 😢

Comment From: sdeleuze

CoWebFilter now allows to customize, potentially dynamically based on the request, the CoroutinesContext.

Comment From: tekener

Looks like with this implementation always the coroutine context of the last filter is winning. I was expecting that you could add more context information to the context created by the previous filter.

Comment From: sdeleuze

Please create a new dedicated issue with a reproducer if you think there is something to refine/fix.