Consistent handling of coroutine context

[ilya40umov]

While refactoring the filter chain used in one of our services, which is based on Kotlin, Spring Boot, WebFlux, coRouter & coroutines, I've run in the following scenario:

• observationRegistry.asContextElement() needs to be added early on to the context, so that the observation from http request is correctly propagated
• multiple separate concerns, such as adding trace baggage, logging incoming request etc., need to be implemented as individual filters
• since we are using coRouter, some of the filters need to be only applied to some of the routes defined in the DSL

Here are the facilities I'm aware of / was able to find, which seem to be relevant for the problem at hand:

• CoWebFilter
• fun filter(filterFunction: suspend (ServerRequest, suspend (ServerRequest) -> ServerResponse) -> ServerResponse) in CoRouterFunctionDsl
• fun context(provider: suspend (ServerRequest) -> CoroutineContext) in CoRouterFunctionDsl

Now, here are the problems I've run into:

• Building a chain of CoWebFilters would require making them all aware of which particular EPs to wrap and which to pass on
• Using fun filter(filterFunction) from CoRouterFunctionDsl allows to apply these in some parts of coRouter, but these filter functions aren't picking up the context that CoWebFilter may have left in COROUTINE_CONTEXT_ATTRIBUTE.
• Additionally, all filters created by fun filter(filterFunction) will not inherit context from one another and aren't able to modify the context that the actual handler will use
• If fun context(provider) is used, it's executed multiple times for 1 request. I.e. it will be called to create a context for each filter, and then for the corresponding handler.

In the end, I've ended up with the following "magical" implementation:

coRouter {
    context { request ->
        if (CoWebFilter.COROUTINE_CONTEXT_ATTRIBUTE !in request.attributes()) {
            request.attributes()[CoWebFilter.COROUTINE_CONTEXT_ATTRIBUTE] =
                Dispatchers.Unconfined + observationRegistry.asContextElement()
        }
        request.attributes()[CoWebFilter.COROUTINE_CONTEXT_ATTRIBUTE] as CoroutineContext
    }
    filter(baggageAddingFilter)
    filter(requestLoggingFilter)
    routes()
}

which is at least able to meet our current needs, but it still has a problem that filters added this way would only be able to modify coroutineContext of one another by modifying request.attributes()[CoWebFilter.COROUTINE_CONTEXT_ATTRIBUTE] explicitly.

IMO, Spring Framework could:

• provide more consistent support for persisting/inheriting coroutine context between parts of the execution chain.
• potentially look into adding a facility similar to fun context(provider) of coRouter that would be executed early on and provide the context for the first `CoWebFilter in the chain
• reevaluate how many times fun context(provider) should be executed by coRouter during handling of a single request (e.g. it could be for example used as a fallback to provide coroutineContext once, if by the time execution goes into coRouter code there was no CoWebFilter invoked).

Tested on: Spring Boot 3.2.4 / Spring 6.1.5