Regression in 3.0.2: tracer context lost when used with spring-data-mongodb

Describe the bug

Hi, I think there’s a regression introduced in 3.0.2. Given an endpoint like:

import org.springframework.cloud.sleuth.Tracer

@GetMapping("chain-with-mongo")
fun chainMongo(): Mono<SpanWrapper> {
    val result = SpanWrapper()
    return tracer.spanId()
           .map { result.add(it) }
           .flatMap { repo.insert(Person(it)) }
           .flatMap { tracer.spanId() } // <-- spanId is empty here
           .map { result.add(it) }
           .map { result }
}

fun Tracer.spanId() = currentSpan()?.context()?.spanId().orEmpty().toMono()

After the reactive repository, traceId or spanId are empty.
This happens from 3.0.2. Previous versions (3.0.1 & 3.0.0) are fine.
Other versions:

• Spring Cloud 2020.0.2
• Spring Boot 2.4.4

If I were to guess where this might come from I would point at the recent optimizations in Reactor operator decoration.

Sample

You can find a reproducer here. You can sample the issue by executing ./gradlew test

[marcingrzejszczak]

BTW - if you bump mongodb to version <mongodb.version>4.2.3</mongodb.version> then things work out of the box without any changes! That’s beacuse Mongo started using project reactor under the hood. I can close this issue then since the recommended approach is to use version 4.2 of the mongodb driver

[marcingrzejszczak]

Hi! The fact that it used to work had to be accidental. We have done a big analysis with @mp911de and we need to talk to Mongo folks about providing hook points to allow the reactive context propagation. Until now as a workaround you have to IMO use the WebFluxSleuthOperators and the like:

@RestController
@RequestMapping("/spans")
public class TestController {

    private static final Logger log = LoggerFactory.getLogger(TestController.class);
    
    private final Tracer tracer;
    private final CurrentTraceContext currentTraceContext;
    private final PersonRepository repo;

    public TestController(Tracer tracer, CurrentTraceContext currentTraceContext, PersonRepository repo) {
        this.tracer = tracer;
        this.currentTraceContext = currentTraceContext;
        this.repo = repo;
    }
    
    @GetMapping("chain-with-mongo")
    Mono<SpanWrapper> chainMongo(ServerWebExchange serverWebExchange) {
        SpanWrapper result = new SpanWrapper().setSpans(new ArrayList<>());
        return spanId(tracer)
                .map(result::add)
                .doOnNext(s -> log.info("ASD1: [{}]", s))
                .flatMap(s -> repo.insert(new Person(s)))
// log stuff with tracing context
                .doOnNext(person -> WebFluxSleuthOperators.withSpanInScope(tracer, currentTraceContext, serverWebExchange, () -> log.info("ASD 2")))
// retrieve the root span and put it in thread local (the scope should be cleared by the TraceFilter later on)
                .doOnNext(person -> tracer.withSpan(serverWebExchange.getAttribute(Span.class.getName())))
                .flatMap(p -> spanId(tracer)) // <-- spanId is empty here
                .map(result::add)
                .map(it -> result);
    }
}

I’ll mark this as in progress cause we can’t do much about it ATM.