Getting reference to context in Spring Reactor - spring

Question

I am using Spring projectreactor reactor-core 3.1.8.RELEASE. I am implementing a logging framework for my microservice to have JSON Audit logs, so used context to store certain fields such as userID, collaboration ID, component Name and few other fields that are common across request life-cycle. Since Threadlocal cannot be used in reactive services to stores these elements, I have to use the context. But getting a reference to the context is apparently very difficult. I can get a reference to the context from the Signal through the doOnEach function call and that's it. If I use doOnEach, it gets called for all signals types and I am unable to isolate on Error, success and so on. Moreover, if an error occurs in between, then all the subsequent doOnEach gets called anyway, so the logs get repeated with several onError log types.
There is very limited documentation regarding how to get a reference to the context object in Spring reactor. Any help regarding a better way to generate audit logs that has collaboration IDs and other request specific IDs stored and propagated across function calls and external invocations is appreciated.
Code Snippets -
In the WebFilter, I am setting few key-value pairs as follows -
override fun filter(exchange: ServerWebExchange, filterChain: WebFilterChain): Mono<Void> {
// add the context variables at the end of the chain as the context moves from
// downstream to upstream.
return filterChain.filter(exchange)
.subscriberContext { context ->
var ctx = context.put(RestRequestInfo::class.java, restRequestInfo(exchange))
ctx = ctx.put(COLLABORATION_ID, UUID.randomUUID().toString())
ctx=ctx.put(COMPONENT_NAME, "sample-component-name")
ctx=ctx.put(USER_NAME, "POSTMAN")
ctx
}
}
Then I want to use the key-value pairs added above in all the subsequent logs so that log aggregators like Splunk can get all the JSON logs associated with this particular request, based on collaboration ID. Right now, the only way to get values out of context is through doOnEach function call, where we get a handle to the SIgnal through which we get handle to context. But all doOnEach gets called during each and every events, irrespective of whether each function call was success or failure
return Mono.just(request)
.doOnEach(**Code to log with context data**)
.map(RequestValidations::validateRequest)
.doOnEach(**Code to log with context data**)
.map(RequestValidations::buildRequest)
.map(RequestValidations::validateQueryParameters)
.doOnEach(**Code to log with context data**)
.flatMap(coverageSummariesGateway::getCoverageSummaries)
.doOnEach(**Code to log with context data**)
.map({ coverageSummaries ->
getCoverageSummariesResponse(coverageSummaries, serviceReferenceId) })
.doOnEach(**Code to log with context data**)
.flatMap(this::renderSuccess)
.doOnEach(**Code to log with context data**)
.doOnError { logger.info("ERROR OCCURRED") }
Thank you!

Answer 1

You could do something like following:
return Mono.just(request)
.doOnEach(**Code to log with context data**)
.flatMap( r -> withMDC(r, RequestValidations::validateRequest))
following method will populate mapping diagnostic context (MDC) so you have it automatically in your logs (depends on you logging pattern). E.g. logback has %X{traceId} where traceId is a key in the tracingContext map.
public static <T, R> Mono<R> withMDC(T value, Function<T, Mono<R>> f) {
return Mono.subscriberContext()
.flatMap( ctx -> {
Optional<Map> tracingContext = ctx.getOrEmpty("tracing-context-key");
if (tracingContext.isPresent()) {
try {
MDC.setContextMap(tracingContext.get());
return f.apply(value);
} finally {
MDC.clear();
}
} else{
return f.apply(value);
}
});
}
It is not quite nice, hope it will be eventually improved by Logging frameworks and context will be injected automatically.