Parent traceId is not reused when calling WebClient.awaitExchange function

[grassehh]

Hi, we are migrating our project from Spring boot 3.5 to Spring Boot 4 and encountered an issue with observability.
Here is a simplified version of one of our tests:

@SpringBootTest
@AutoConfigureTracing
@ActiveProfiles("test")
@EnableWireMock
class TracingTest @Autowired constructor(
    private val observationRegistry: ObservationRegistry,
    private val webClient: WebClient,
    private val wireMock: WireMock,
    private val tracer: Tracer,
    @param:Value($$"${wiremock.server.port}") private val wireMockServerPort: Int
) {
    @Test
    fun test() {
        var traceId: String? = ""
        Observation.createNotStarted("test", observationRegistry).observe {
            traceId = tracer.currentTraceContext().context()?.traceId()
            runBlocking {
                webClient.get()
                    .uri("http://localhost:${wireMockServerPort}/test")
                    .awaitExchange { }
                webClient.get()
                    .uri("http://localhost:${wireMockServerPort}/test")
                    .awaitExchange { }
            }
        }

        await untilAsserted {
            wireMock.verifyThat(
                2,
                getRequestedFor(urlEqualTo("/test"))
                    .withHeader("X-B3-TraceId", equalTo(traceId))
            )
        }
    }
}

It creates an observation, executes a webclient call twice inside it, awaits the exchange inside, then checks that the traceId is the same for both calls (through the X-B3-TraceId) propagated header.

With Spring Boot 4, the test fails, but if we revert the code of the awaitExchange function as it was before this commit it works again. So basically, using this awaitExchange implementation:

suspend fun <T: Any> WebClient.RequestHeadersSpec<out WebClient.RequestHeadersSpec<*>>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T {
        val context = currentCoroutineContext().minusKey(Job.Key)
        return exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingle()
    }

Is there anything we were doing wrong and it was not supposed to work or is there a regression ?
Thanks.

[grassehh]

Hello, any ETA about any idea ? This is the last subject that's preventing us to merge our Spring Boot 4 migration :(
At least maybe having a workaround :)
Thank you !

[sdeleuze]

Hey, sorry for the delay, it was in my todo but busy times.

There may be a need for a refinement to merge instead of replacing the context, but not sure yet cc @dmitrysulman @bclozel in case they have a take on this.

Could you please provide a reproducer as a Spring Boot application in addition to the test mentioned in the description?

Do you leverage automatic context propagation with coroutines (see related Framework documentation)?

[grassehh]

Hello @sdeleuze, no problem, thanks for the quick reply :)
I'm still trying to create a reproducer, but for some reasons I cannot reproduce on the sample.
Ironically though, on our project, removing context-propagation: auto property makes it work again... Maybe that's a clue. I'll try to dig more.

[grassehh]

Ok I managed to get a reproducer.
Checkout this repo https://github.com/grassehh/issues then run the ApplicationTest test class inside spring-36182 project.
The test will fail.
However If you either remove context-propagation: auto property or uncomment the private awaitExchange function in ExampleService, the test will pass.

[dmitrysulman]

Thanks for the reproducer. I debugged it and I have some thoughts on why this is happening.

When reactor context propagation to ThreadLocals is enabled, every contextWrite(...) call updates the ThreadLocal value with whatever context is passed to the contextWrite method. The new awaitExchange implementation introduced in #34555 leverages the ReactorContext coroutine context element to propagate the coroutine context to downstream exchange filter functions.

The interop between the coroutine context and reactor context is implemented in the kotlinx-coroutines-reactor module. The logic is to extract the ReactorContext element from the coroutine context and pass it to the Reactor Publisher using contextWrite(...) (see the ReactorContextInjector implementation). This contextWrite invocation triggers a propagation to the ThreadLocal used by Micrometer, but the context doesn't contain an observation at this step, so the tracing is lost.

The workaround to make the provided test work is to refine the CoroutineContext.toReactorContext() function to capture the Micrometer context snapshot and add it to the same ReactorContext coroutine context element:

    private val contextSnapshotFactory = ContextSnapshotFactory.builder()
        .contextRegistry(ContextRegistry.getInstance()).build()

    @PublishedApi
    internal fun CoroutineContext.toReactorContext(): ReactorContext {
        val threadLocalContextSnapshot = contextSnapshotFactory.captureAll()
        val context = Context.of(COROUTINE_CONTEXT_ATTRIBUTE, this).let {
            threadLocalContextSnapshot.updateContext(it)
        }.readOnly()
        return (this[ReactorContext.Key]?.context?.putAll(context) ?: context).asCoroutineContext()
    }

Note that currently, this code can't be compiled as-is because the spring-webflux module doesn't have the context-propagation library on the classpath. If we decide to go this way, a possible solution is to add a utility function to spring-core that can enrich any Reactor context with a Micrometer context snapshot. I can work on this, but I will need some guidance on how to better organize the spring-core part.

[grassehh]

Thanks @dmitrysulman
So if I understand, there is no possible workaround on my side at the moment ?
What I don't understand is that tracing propagation works well on Spring-instrumented observations (like for rest controllers or AMQP listeners). However in our case, we are creating the observation ourselves from a MongoChangeStream.

[sdeleuze]

FYI I will be on PTO for 2 weeks, feel free to explore the topic, I will decide what to do when I am back, so any exploration/discussion will help. Thanks for your patience.

[be-hase]

Hello. I’ve run into this issue as well and it’s been a real pain.

I know it can be worked around by writing bodyToMono(...).awaitSingle() without using this extension function
, but I’d really appreciate it if the behavior could be reverted to what it was before.

Additional note:

Personally, I think it would be less confusing if awaitBody() simply stuck to being shorthand for bodyToMono<>().awaitSingle().
It feels like this issue might have been better fixed with a different approach… (though I can’t think of what that approach should be…).

[grassehh]

Hi, any update please ? 🙏