Propagate CoroutineContext in CoWebFilter

[be-hase]

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

[be-hase]

To give you another example, grpc-kotlin supports customizing coroutine context.
grpc/grpc-kotlin#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.

[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?

[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.

[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.

[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.

spring-framework/spring-webflux/src/main/java/org/springframework/web/reactive/result/method/InvocableHandlerMethod.java

Lines 139 to 145 in 983c6e2

	Method method = getBridgedMethod();
	if (KotlinDetector.isSuspendingFunction(method)) {
	value = CoroutinesUtils.invokeSuspendingFunction(method, getBean(), args);
	}
	else {
	value = method.invoke(getBean(), args);
	}

[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.

[be-hase]

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

[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

[sdeleuze]

Let's try to see if we can leverage CoWebFilter to support this use case by taking inspiration of #26977 (comment).

[be-hase]

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

[sdeleuze]

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

[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.

[sdeleuze]

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