Implement trace-id with Spring Webflux - spring-boot

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
client -> service1 - service1 should generate traceId and log request
service1 -> service2 - service2 should fetch traceId from request, then log request
service1 <- service2 - after some calculation service2 should log response and return response to service1
client <- service1 - at the end service1 should log response (still with the same traceId) and return response to client
How it actually works
client -> service1 - nothing in logs
service1 -> service2 - nothings in logs
service1 <- service2 - service2 is logging correctly and return response to service1
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

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);
.....
)

Answer 2

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