Implement trace-id with Spring Webflux

Question

I'd like to generate unique traceId per request and pass it through all services. In Spring MVC it was fairly easy by using MDC context and putting traceId in header, but in reactive stack it isn't working at all because of ThreadLocal.

In general I'd like to log each request and response on every service I have with single traceId which can identify specific action in whole system.

I tried to create custom filter based on article: https://azizulhaq-ananto.medium.com/how-to-handle-logs-and-tracing-in-spring-webflux-and-microservices-a0b45adc4610 but it's seems to not working. My current solution only log responses and traceId are losing after making request, so there is no on response. Let's try imagine that there are two services: service1 and service2. Below I tried to sketch how it should work.

How should it work

1. client -> service1 - service1 should generate traceId and log request
2. service1 -> service2 - service2 should fetch traceId from request, then log request
3. service1 <- service2 - after some calculation service2 should log response and return response to service1
4. client <- service1 - at the end service1 should log response (still with the same traceId) and return response to client

How it actually works

1. client -> service1 - nothing in logs
2. service1 -> service2 - nothings in logs
3. service1 <- service2 - service2 is logging correctly and return response to service1
4. client <- service1 - service1 is logging response (but without traceId)

Here is my approach

@Component
public class TraceIdFilter implements WebFilter {

    private static final Logger log = LoggerFactory.getLogger(TraceIdFilter.class);

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        Map<String, String> headers = exchange.getRequest().getHeaders().toSingleValueMap();
        return Mono.fromCallable(() ->  {
            final long startTime = System.currentTimeMillis();

            return new ServerWebExchangeDecorator(exchange) {
                @Override
                public ServerHttpRequest getRequest() {
                    return new RequestLoggingInterceptor(super.getRequest(), false);
                }

                @Override
                public ServerHttpResponse getResponse() {
                    return new ResponseLoggingInterceptor(super.getResponse(), startTime, false);
                }
            };
        }).contextWrite(context -> {
            var traceId = "";
            if (headers.containsKey("X-B3-TRACEID")) {
                traceId = headers.get("X-B3-TRACEID");
                MDC.put("X-B3-TraceId", traceId);
            } else if (!exchange.getRequest().getURI().getPath().contains("/actuator")) {
                traceId = UUID.randomUUID().toString();
                MDC.put("X-B3-TraceId", traceId);
            }

            Context contextTmp = context.put("X-B3-TraceId", traceId);
            exchange.getAttributes().put("X-B3-TraceId", traceId);


            return contextTmp;
        }).flatMap(chain::filter);

    }


}

Github: https://github.com/Faelivrinx/kotlin-spring-boot

There is any existing solution do that?

Answer 1

Sleuth 3.0 provides automatic instrumentation for WebFlux, which means, if you do nothing, you will always get the current span of the thread which subscribes to your Mono or Flux. If you wish to overwrite this, for instance because you are batching a lot of requests and want a unique trace for every transaction, all you have to do is to manipulate the Context of your operator chain.

private Mono<Data> performRequest() {
    var span = tracer.spanBuilder().setNoParent().start(); // generate a completely new trace
                                                           // note: you can also just generate a new span if you want
    return Mono.defer(() -> callRealService())
        .contextWrite(Context.of(TraceContext.class, span.context());
}

When following this approach, make sure to import org.springframework.cloud.sleuth.Tracer and not the one by brave, because they use different types and Reactor will drop your Mono with an ugly error message if they don't align correctly (unfortunately, since the Context ist just a plain old Map<Object, Object>, you won't get a compiler error).

Answer 2

In Spring Webflux you don't have anymore a ThreadLocal but you have a unique context for each chain request. You can attach a traceId to this context as following:

@Component
public class TraceIdFilter implements WebFilter {

    private static final Logger log = LoggerFactory.getLogger(TraceIdFilter.class);

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {

        return chain.filter(exchange)
                .subscriberContext(
                        ctx -> {
                            .....
                            var traceId = UUID.randomUUID().toString();
                            return   ctx.put("X-B3-TraceId", traceId);
                            .....
                        }
                );

    }


}

Now the chain in your service will have in the context this attribute. You can retrive it from your service using static method Mono.subscriberContext(). For example you can get traceId in this way

Mono.subscriberContext()
    .flaMap(ctx -> {
       .....
       var traceId = ctx.getOrDefault("traceId", null);
       .....
    )