Reactive Programming: Spring WebFlux: How to build a chain of micro-service calls? - spring-boot

Spring Boot Application:
a #RestController receives the following payload:
{
"cartoon": "The Little Mermaid",
"characterNames": ["Ariel", "Prince Eric", "Sebastian", "Flounder"]
}
I need to process it in the following way:
Get the unique Id for each character name: make an HTTP call to "cartoon-characters" microservice, that returns ids by names
Transform data received by the controller:
replace character names with appropriate ids that were received on the previous step from "cartoon-characters" microservice.
{
"cartoon": "The Little Mermaid",
"characterIds": [1, 2, 3, 4]
}
Send an HTTP POST request to "cartoon-db" microservice with transformed data.
Map the response from "cartoon-db" to the internal representation that is the controller return value.
The problem that I got:
I need to implement all these steps using the paradigm of Reactive Programming (non-blocking\async processing) with Spring WebFlux (Mono|Flux) and Spring Reactive WebClient - but I have zero experience with that stack, trying to read about it as much as I can, plus googling a lot but still, have a bunch of unanswered questions, for example:
Q1. I have already configured reactive webClient that sends a request to "cartoon-characters" microservice:
public Mono<Integer> getCartoonCharacterIdbyName(String characterName) {
return WebClient.builder().baseUrl("http://cartoon-characters").build()
.get()
.uri("/character/{characterName}", characterName)
.retrieve()
.bodyToMono(Integer.class);
}
As you may see, I have got a list of cartoon character names and for each of them I need to call getCartoonCharacterIdbyName(String name) method, I am not sure that the right option to call it in series, believe the right option: parallel execution.
Wrote the following method:
public List<Integer> getCartoonCharacterIds(List<String> names) {
Flux<Integer> flux = Flux.fromStream(names.stream())
.flatMap(this::getCartoonCharacterIdbyName);
return StreamSupport.stream(flux.toIterable().spliterator(), false)
.collect(Collectors.toList());
}
but I have doubts, that this code does parallel WebClient execution and also, code calls flux.toIterable() that block the thread, so with this implementation I lost non-blocking mechanism.
Are my assumptions correct?
How do I need to rewrite it to having parallelism and non-blocking?
Q2.
Is it technically possible to transform input data received by the controller (I mean replace names with ids) in reactive style: when we operate with Flux<Integer> characterIds, but not with the List<Integer> of characterIds?
Q3. Is it potentially possible to get not just transformed Data object, but Mono<> after step 2 that can be consumed by another WebClient in Step 3?

Actually it's a good question since understanding the WebFlux, or project reactor framework, when it comes to chaining micro-services requires a couple of steps.
The first is to realize that a WebClient should take a publisher in and return a publisher. Extrapolate this to 4 different method signatures to help with thinking.
Mono -> Mono
Flux -> Flux
Mono -> Flux
Flux -> Mono
For sure, in all cases, it is just Publisher->Publisher, but leave that until you understand things better. The first two are obvious, and you just use .map(...) to handle objects in the flow, but you need to learn how to handle the second two. As commented above, going from Flux->Mono could be done with .collectList(), or also with .reduce(...). Going from Mono->Flux seems to generally be done with .flatMapMany or .flatMapIterable or some variation of that. There are probably other techniques. You should never use .block() in any WebFlux code, and generally you will get a runtime error if you try to do so.
In your example you want to go to
(Mono->Flux)->(Flux->Flux)->(Flux->Flux)
As you said, you want
Mono->Flux->Flux
The second part is to understand about chaining Flows. You could do
p3(p2(p1(object)));
Which would chain p1->p2->p3, but I always found it more understandable to make a "Service Layer" instead.
o2 = p1(object);
o3 = p2(o2);
result = p3(o3);
This code is just much easier to read and maintain and, with some maturity, you come to understand the worth of that statement.
The only problem I had with your example was doing a Flux<String> with WebClient as a #RequestBody. Doesn't work. See WebClient bodyToFlux(String.class) for string list doesn't separate individual values. Other than that, it's a pretty straightforward application. You'll find when you debug it that it gets to the .subscribe(System.out::println) line before it gets to the Flux<Integer> ids = mapNamesToIds(fn) line. This is because the Flow is setup before it is executed. Takes a while to understand this but it is the point of the project reactor framework.
#SpringBootApplication
#RestController
#RequestMapping("/demo")
public class DemoApplication implements ApplicationRunner {
public static void main(String[] args) {
SpringApplication.run(DemoApplication.class, args);
}
Map<Integer, CartoonCharacter> characters;
#Override
public void run(ApplicationArguments args) throws Exception {
String[] names = new String[] {"Ariel", "Prince Eric", "Sebastian", "Flounder"};
characters = Arrays.asList( new CartoonCharacter[] {
new CartoonCharacter(names[0].hashCode(), names[0], "Mermaid"),
new CartoonCharacter(names[1].hashCode(), names[1], "Human"),
new CartoonCharacter(names[2].hashCode(), names[2], "Crustacean"),
new CartoonCharacter(names[3].hashCode(), names[3], "Fish")}
)
.stream().collect(Collectors.toMap(CartoonCharacter::getId, Function.identity()));
// TODO Auto-generated method stub
CartoonRequest cr = CartoonRequest.builder()
.cartoon("The Little Mermaid")
.characterNames(Arrays.asList(names))
.build();
thisLocalClient
.post()
.uri("cartoonDetails")
.body(Mono.just(cr), CartoonRequest.class)
.retrieve()
.bodyToFlux(CartoonCharacter.class)
.subscribe(System.out::println);
}
#Bean
WebClient localClient() {
return WebClient.create("http://localhost:8080/demo/");
}
#Autowired
WebClient thisLocalClient;
#PostMapping("cartoonDetails")
Flux<CartoonCharacter> getDetails(#RequestBody Mono<CartoonRequest> cartoonRequest) {
Flux<StringWrapper> fn = cartoonRequest.flatMapIterable(cr->cr.getCharacterNames().stream().map(StringWrapper::new).collect(Collectors.toList()));
Flux<Integer> ids = mapNamesToIds(fn);
Flux<CartoonCharacter> details = mapIdsToDetails(ids);
return details;
}
// Service Layer Methods
private Flux<Integer> mapNamesToIds(Flux<StringWrapper> names) {
return thisLocalClient
.post()
.uri("findIds")
.body(names, StringWrapper.class)
.retrieve()
.bodyToFlux(Integer.class);
}
private Flux<CartoonCharacter> mapIdsToDetails(Flux<Integer> ids) {
return thisLocalClient
.post()
.uri("findDetails")
.body(ids, Integer.class)
.retrieve()
.bodyToFlux(CartoonCharacter.class);
}
// Services
#PostMapping("findIds")
Flux<Integer> getIds(#RequestBody Flux<StringWrapper> names) {
return names.map(name->name.getString().hashCode());
}
#PostMapping("findDetails")
Flux<CartoonCharacter> getDetails(#RequestBody Flux<Integer> ids) {
return ids.map(characters::get);
}
}
Also:
#Data
#NoArgsConstructor
#AllArgsConstructor
#Builder
public class StringWrapper {
private String string;
}
#Data
#Builder
public class CartoonRequest {
private String cartoon;
private List<String> characterNames;
}
#Data
#Builder
#NoArgsConstructor
#AllArgsConstructor
public class CartoonCharacter {
Integer id;
String name;
String species;
}

Related

Spring Integration Flow with #Restcontoller Timing issue

A simple #RestController is connected with a #MessagingGateway to an IntegrationFlow.
After a load test we saw within the tracing that we lose "a lot of time" before even starting the processing within the flow:
Tracing result
In this example we can see that over 90ms spend befor sending the message to the flow.
Did anyone have some idea what leads to this behavior?
As far as I understood the documentation, everything is handled in the sender thread and therefore no special worker threads are created.
We use the Restcontroller since we need to create the documentation with springdoc-openapi-ui
ExampleCode:
RestController
#RestController
public class DescriptionEndpoint {
HttpMessageGateway httpMessageGateway;
public Result findData(#Valid dataRequest dataRequest) {
final Map<String, Object> headerParams = new HashMap<>();
return httpMessageGateway.basicDataDescriptionFlow(dataRequest, headerParams);
}
}
Gateway
#MessagingGateway
public interface HttpMessageGateway {
#Gateway(requestChannel = "startDataFlow.input")
Result basicDataDescriptionFlow(#Payload dataRequest prDataRequest, #Headers Map<String, Object> map);
}
IntegrationFlow
public class ExampleFlow {
#Bean
public IntegrationFlow startDataFlow() {
return new FlowExtension()
.handle(someHandler1)
.handle(someHandler2)
.handle(someHandler3)
.get();
}
}
After adding some more traces I realized, that this timing issue is caused by my spring security configuration.
Unfortunatelly, i thought, the span is only representing the time after the start of findData(..). But it seems, the tracing starts already in the proxy methods and security chain.
After improving some implementation on our JWTToken filter, the spend times for these endpoints are OK.

Use JSON transformer in Spring Integration

I have a problem that seems unaddressed in any of the examples I can find.
My application reads an ActiveMQ topic of JSON messages. It will build a completely new outbound REST call based on this data. Note that this is not a "transformation". It is given "X" produce "Y" i.e. ServiceActivator.
My flows thus far are
public IntegrationFlow splitInputFlow() {
return IntegrationFlows.from("inboundJmsChannel")
.split()
.log(LoggingHandler.Level.DEBUG)
.route(Message.class, m -> m.getHeaders().get("x-bn-class").equals("Healthcheck.class") ? "healthcheckChannel" : "metricChannel")
.get();
}
public IntegrationFlow healthcheckFlow() {
return IntegrationFlows.from("healthcheckChannel")
.log(LoggingHandler.Level.DEBUG)
.transform(Transformers.fromJson(Healthcheck.class))
.handle("healthcheckActivator", "process")
.get();
}
There are dozens of examples on how to use spring transformers. I have even considered trying a MessageConverter. But I don't see why it would help and it doesn't seem to be the normal approach.
The main problem here is that Integration calls healthcheckActivator.process(String payload). The payload itself is the expected valid JSON string.
I am a little surprised it does not call healtcheckActivator.process(Message payload) but But that wouldn't help so it doesn't much matter.
The real question is why does it not call healtcheckActivator.process(Healthcheck healthcheck)?
Well actually I understand "why". It is because DSL generates an internal channel to tie the steps together and as far as I understand anything on a channel is a spring.messaging.Message.
I can easily instantiate my Healthcheck object once I get inside the SA. But that leaves the nagging question: What possible good is the entire transform step? If it always "serializes" the object back into a Message -- what's the point.
Like I said I think I'm missing something fundamental here.
EDIT
My new (and probably last) idea is maybe I'm publishing it wrong.
To publish it I am using
jmsTemplate.convertAndSend(topicName, healthcheck, messagePostProcessor -> {
messagePostProcessor.setJMSType("TextMessage");
messagePostProcessor.setStringProperty("x-bn-class", "Healthcheck.class");
messagePostProcessor.setStringProperty("x-bn-service-name", restEndpoint.getServiceName());
messagePostProcessor.setStringProperty("x-bn-service-endpoint-name", restEndpoint.getEndpointName());
messagePostProcessor.setLongProperty("x-bn-heathcheck-timestamp", queryDate);
messagePostProcessor.setStringProperty("x-bn-healthcheck-status", subsystemStatus.getStatus(subsystemStatus));
messagePostProcessor.setIntProperty("httpStatus", httpStatus.value());
return messagePostProcessor;
});
What arrives in the SI process(String payload) method is:
LoggingHandler - GenericMessage [payload={"healthcheckType":"LOCAL","outcome":"PASS","dependencyType":"DB","endpoint":"NODE TABLE","description":"Read from DB","durationSecs":0.025}, headers={x-bn-service-name=TG10-CS2, x-bn-service-endpoint-name=TG Q10-CS2 Ready Check, jms_destination=topic://HEALTH_MONITOR, _type=com.healthcheck.response.Healthcheck, x-bn-heathcheck-timestamp=1558356538000, priority=4, jms_timestamp=1558356544244, x-bn-healthcheck-status=SEV0, jms_redelivered=false, x-bn-class=Healthcheck.class, httpStatus=200, jms_type=TextMessage, id=b29ffea7-7128-c543-9a14-8bab450f0ac6, jms_messageId=ID:39479-1558356520091-1:2:1:1:1, timestamp=1558356544409}]
I hadn't noticed the _type parameter in the jms_destination header before. But before I started screwing around with this (because it didn't work) that is the correct class name for what the other team provided.
I have not implemented a JMS message converter. But the supplied SimpleMessageConverter seems that it should do exactly what I want.
Your understanding is correct; works fine for me, so something else is going on...
#SpringBootApplication
public class So56169938Application {
public static void main(String[] args) {
SpringApplication.run(So56169938Application.class, args);
}
#Bean
public IntegrationFlow flow() {
return IntegrationFlows.from(() -> "{\"foo\":\"bar\"}", e -> e.poller(Pollers.fixedDelay(5000)))
.transform(Transformers.fromJson(Foo.class))
.handle("myBean", "method")
.get();
}
#Bean
public MyBean myBean() {
return new MyBean();
}
public static class MyBean {
public void method(Foo foo) {
System.out.println(foo);
}
}
public static class Foo {
private String foo;
String getFoo() {
return this.foo;
}
void setFoo(String foo) {
this.foo = foo;
}
#Override
public String toString() {
return "Foo [foo=" + this.foo + "]";
}
}
}
and
Foo [foo=bar]
Foo [foo=bar]
Foo [foo=bar]
Foo [foo=bar]
Foo [foo=bar]
Foo [foo=bar]
Well, Spring Integration is a Messaging framework. It transfers messages from endpoint to endpoint via channels in between. That's already the target endpoint responsibility to deal with consumed message the proper way. The framework doesn't care about the payload. It is really a business part of the target application. That's how we can make framework components as generic as possible leaving the room for target business types for end-users.
Anyway the Framework provides some mechanisms to interact with payloads. We call it POJO method invocation. So, you provide some business with arbitrary contract, however following some Spring Integration rules: https://docs.spring.io/spring-integration/docs/current/reference/html/#service-activator.
So, according your description it is really a surprise that it doesn't work for healtcheckActivator.process(Healthcheck healthcheck). Your transform(Transformers.fromJson(Healthcheck.class)) should really produce a Message with Healthcheck object as a payload. The framework consults a method signature and tries to map a payload and/or headers to the method invocation arguments, having the whole message as a container for data to delegate to the method call.
From here it would be great to see your healtcheckActivator.process() method to determine why the transform(Transformers.fromJson(Healthcheck.class)) result cannon be mapped to that method arguments.

Spring Cloud - HystrixCommand - How to properly enable with shared libraries

Using Springboot 1.5.x, Spring Cloud, and JAX-RS:
I could use a second pair of eyes since it is not clear to me whether the Spring configured, Javanica HystrixCommand works for all use cases or whether I may have an error in my code. Below is an approximation of what I'm doing, the code below will not actually compile.
From below WebService lives in a library with separate package path to the main application(s). Meanwhile MyWebService lives in the application that is in the same context path as the Springboot application. Also MyWebService is functional, no issues there. This just has to do with the visibility of HystrixCommand annotation in regards to Springboot based configuration.
At runtime, what I notice is that when a code like the one below runs, I do see "commandKey=A" in my response. This one I did not quite expect since it's still running while the data is obtained. And since we log the HystrixRequestLog, I also see this command key in my logs.
But all the other Command keys are not visible at all, regardless of where I place them in the file. If I remove CommandKey-A then no commands are visible whatsoever.
Thoughts?
// Example WebService that we use as a shared component for performing a backend call that is the same across different resources
#RequiredArgsConstructor
#Accessors(fluent = true)
#Setter
public abstract class WebService {
private final #Nonnull Supplier<X> backendFactory;
#Setter(AccessLevel.PACKAGE)
private #Nonnull Supplier<BackendComponent> backendComponentSupplier = () -> new BackendComponent();
#GET
#Produces("application/json")
#HystrixCommand(commandKey="A")
public Response mainCall() {
Object obj = new Object();
try {
otherCommandMethod();
} catch (Exception commandException) {
// do nothing (for this example)
}
// get the hystrix request information so that we can determine what was executed
Optional<Collection<HystrixInvokableInfo<?>>> executedCommands = hystrixExecutedCommands();
// set the hystrix data, viewable in the response
obj.setData("hystrix", executedCommands.orElse(Collections.emptyList()));
if(hasError(obj)) {
return Response.serverError()
.entity(obj)
.build();
}
return Response.ok()
.entity(healthObject)
.build();
}
#HystrixCommand(commandKey="B")
private void otherCommandMethod() {
backendComponentSupplier
.get()
.observe()
.toBlocking()
.subscribe();
}
Optional<Collection<HystrixInvokableInfo<?>>> hystrixExecutedCommands() {
Optional<HystrixRequestLog> hystrixRequest = Optional
.ofNullable(HystrixRequestLog.getCurrentRequest());
// get the hystrix executed commands
Optional<Collection<HystrixInvokableInfo<?>>> executedCommands = Optional.empty();
if (hystrixRequest.isPresent()) {
executedCommands = Optional.of(hystrixRequest.get()
.getAllExecutedCommands());
}
return executedCommands;
}
#Setter
#RequiredArgsConstructor
public class BackendComponent implements ObservableCommand<Void> {
#Override
#HystrixCommand(commandKey="Y")
public Observable<Void> observe() {
// make some backend call
return backendFactory.get()
.observe();
}
}
}
// then later this component gets configured in the specific applications with sample configuraiton that looks like this:
#SuppressWarnings({ "unchecked", "rawtypes" })
#Path("resource/somepath")
#Component
public class MyWebService extends WebService {
#Inject
public MyWebService(Supplier<X> backendSupplier) {
super((Supplier)backendSupplier);
}
}
There is an issue with mainCall() calling otherCommandMethod(). Methods with #HystrixCommand can not be called from within the same class.
As discussed in the answers to this question this is a limitation of Spring's AOP.

Spring Webflux and #Cacheable - proper way of caching result of Mono / Flux type

I'm learning Spring WebFlux and during writing a sample application I found a concern related to Reactive types (Mono/Flux) combined with Spring Cache.
Consider the following code-snippet (in Kotlin):
#Repository
interface TaskRepository : ReactiveMongoRepository<Task, String>
#Service
class TaskService(val taskRepository: TaskRepository) {
#Cacheable("tasks")
fun get(id: String): Mono<Task> = taskRepository.findById(id)
}
Is this valid and safe way of caching method calls returning Mono or Flux? Maybe there are some other principles to do this?
The following code is working with SimpleCacheResolver but by default fails with Redis because of the fact that Mono is not Serializable. In order to make them work e.g Kryo serializer needs to be used.
Hack way
For now, there is no fluent integration of #Cacheable with Reactor 3.
However, you may bypass that thing by adding .cache() operator to returned Mono
#Repository
interface TaskRepository : ReactiveMongoRepository<Task, String>
#Service
class TaskService(val taskRepository: TaskRepository) {
#Cacheable("tasks")
fun get(id: String): Mono<Task> = taskRepository.findById(id).cache()
}
That hack cache and share returned from taskRepository data. In turn, spring cacheable will cache a reference of returned Mono and then, will return that reference. In other words, it is a cache of mono which holds the cache :).
Reactor Addons Way
There is an addition to Reactor 3 which allows fluent integration with modern in-memory caches like caffeine, jcache, etc. Using that technique you will be capable to cache your data easily:
#Repository
interface TaskRepository : ReactiveMongoRepository<Task, String>
#Service
class TaskService(val taskRepository: TaskRepository) {
#Autowire
CacheManager manager;
fun get(id: String): Mono<Task> = CacheMono.lookup(reader(), id)
.onCacheMissResume(() -> taskRepository.findById(id))
.andWriteWith(writer());
fun reader(): CacheMono.MonoCacheReader<String, Task> = key -> Mono.<Signal<Task>>justOrEmpty((Signal) manager.getCache("tasks").get(key).get())
fun writer(): CacheMono.MonoCacheWriter<String, Task> = (key, value) -> Mono.fromRunnable(() -> manager.getCache("tasks").put(key, value));
}
Note: Reactor addons caching own abstraction which is Signal<T>, so, do not worry about that and following that convention
I have used Oleh Dokuka's hacky solution worked great but there is a catch. You must use a greater Duration in Flux cache than your Cachable caches timetolive value. If you dont use a duration for Flux cache it wont invalidate it (Flux documentation says "Turn this Flux into a hot source and cache last emitted signals for further Subscriber.").
So making Flux cache 2 minutes and timetolive 30 seconds can be valid configuration. If ehcahce timeout occurs first, than a new Flux cache reference is generated and it will be used.
// In a Facade:
public Mono<HybrisResponse> getProducts(HybrisRequest request) {
return Mono.just(HybrisResponse.builder().build());
}
// In a service layer:
#Cacheable(cacheNames = "embarkations")
public HybrisResponse cacheable(HybrisRequest request) {
LOGGER.info("executing cacheable");
return null;
}
#CachePut(cacheNames = "embarkations")
public HybrisResponse cachePut(HybrisRequest request) {
LOGGER.info("executing cachePut");
return hybrisFacade.getProducts(request).block();
}
// In a Controller:
HybrisResponse hybrisResponse = null;
try {
// get from cache
hybrisResponse = productFeederService.cacheable(request);
} catch (Throwable e) {
// if not in cache then cache it
hybrisResponse = productFeederService.cachePut(request);
}
return Mono.just(hybrisResponse)
.map(result -> ResponseBody.<HybrisResponse>builder()
.payload(result).build())
.map(ResponseEntity::ok);

Does CompletableFuture have a corresponding Local context?

In the olden days, we had ThreadLocal for programs to carry data along with the request path since all request processing was done on that thread and stuff like Logback used this with MDC.put("requestId", getNewRequestId());
Then Scala and functional programming came along and Futures came along and with them came Local.scala (at least I know the twitter Futures have this class). Future.scala knows about Local.scala and transfers the context through all the map/flatMap, etc. etc. functionality such that I can still do Local.set("requestId", getNewRequestId()); and then downstream after it has travelled over many threads, I can still access it with Local.get(...)
Soooo, my question is in Java, can I do the same thing with the new CompletableFuture somewhere with LocalContext or some object (not sure of the name) and in this way, I can modify Logback MDC context to store it in that context instead of a ThreadLocal such that I don't lose the request id and all my logs across the thenApply, thenAccept, etc. etc. still work just fine with logging and the -XrequestId flag in Logback configuration.
EDIT:
As an example. If you have a request come in and you are using Log4j or Logback, in a filter, you will set MDC.put("requestId", requestId) and then in your app, you will log many log statements line this:
log.info("request came in for url="+url);
log.info("request is complete");
Now, in the log output it will show this:
INFO {time}: requestId425 request came in for url=/mypath
INFO {time}: requestId425 request is complete
This is using a trick of ThreadLocal to achieve this. At Twitter, we use Scala and Twitter Futures in Scala along with a Local.scala class. Local.scala and Future.scala are tied together in that we can achieve the above scenario still which is very nice and all our log statements can log the request id so the developer never has to remember to log the request id and you can trace through a single customers request response cycle with that id.
I don't see this in Java :( which is very unfortunate as there are many use cases for that. Perhaps there is something I am not seeing though?
If you come across this, just poke the thread here
http://mail.openjdk.java.net/pipermail/core-libs-dev/2017-May/047867.html
to implement something like twitter Futures which transfer Locals (Much like ThreadLocal but transfers state).
See the def respond() method in here and how it calls Locals.save() and Locals.restort()
https://github.com/simonratner/twitter-util/blob/master/util-core/src/main/scala/com/twitter/util/Future.scala
If Java Authors would fix this, then the MDC in logback would work across all 3rd party libraries. Until then, IT WILL NOT WORK unless you can change the 3rd party library(doubtful you can do that).
My solution theme would be to (It would work with JDK 9+ as a couple of overridable methods are exposed since that version)
Make the complete ecosystem aware of MDC
And for that, we need to address the following scenarios:
When all do we get new instances of CompletableFuture from within this class? → We need to return a MDC aware version of the same rather.
When all do we get new instances of CompletableFuture from outside this class? → We need to return a MDC aware version of the same rather.
Which executor is used when in CompletableFuture class? → In all circumstances, we need to make sure that all executors are MDC aware
For that, let's create a MDC aware version class of CompletableFuture by extending it. My version of that would look like below
import org.slf4j.MDC;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.function.Supplier;
public class MDCAwareCompletableFuture<T> extends CompletableFuture<T> {
public static final ExecutorService MDC_AWARE_ASYNC_POOL = new MDCAwareForkJoinPool();
#Override
public CompletableFuture newIncompleteFuture() {
return new MDCAwareCompletableFuture();
}
#Override
public Executor defaultExecutor() {
return MDC_AWARE_ASYNC_POOL;
}
public static <T> CompletionStage<T> getMDCAwareCompletionStage(CompletableFuture<T> future) {
return new MDCAwareCompletableFuture<>()
.completeAsync(() -> null)
.thenCombineAsync(future, (aVoid, value) -> value);
}
public static <T> CompletionStage<T> getMDCHandledCompletionStage(CompletableFuture<T> future,
Function<Throwable, T> throwableFunction) {
Map<String, String> contextMap = MDC.getCopyOfContextMap();
return getMDCAwareCompletionStage(future)
.handle((value, throwable) -> {
setMDCContext(contextMap);
if (throwable != null) {
return throwableFunction.apply(throwable);
}
return value;
});
}
}
The MDCAwareForkJoinPool class would look like (have skipped the methods with ForkJoinTask parameters for simplicity)
public class MDCAwareForkJoinPool extends ForkJoinPool {
//Override constructors which you need
#Override
public <T> ForkJoinTask<T> submit(Callable<T> task) {
return super.submit(MDCUtility.wrapWithMdcContext(task));
}
#Override
public <T> ForkJoinTask<T> submit(Runnable task, T result) {
return super.submit(wrapWithMdcContext(task), result);
}
#Override
public ForkJoinTask<?> submit(Runnable task) {
return super.submit(wrapWithMdcContext(task));
}
#Override
public void execute(Runnable task) {
super.execute(wrapWithMdcContext(task));
}
}
The utility methods to wrap would be such as
public static <T> Callable<T> wrapWithMdcContext(Callable<T> task) {
//save the current MDC context
Map<String, String> contextMap = MDC.getCopyOfContextMap();
return () -> {
setMDCContext(contextMap);
try {
return task.call();
} finally {
// once the task is complete, clear MDC
MDC.clear();
}
};
}
public static Runnable wrapWithMdcContext(Runnable task) {
//save the current MDC context
Map<String, String> contextMap = MDC.getCopyOfContextMap();
return () -> {
setMDCContext(contextMap);
try {
return task.run();
} finally {
// once the task is complete, clear MDC
MDC.clear();
}
};
}
public static void setMDCContext(Map<String, String> contextMap) {
MDC.clear();
if (contextMap != null) {
MDC.setContextMap(contextMap);
}
}
Below are some guidelines for usage:
Use the class MDCAwareCompletableFuture rather than the class CompletableFuture.
A couple of methods in the class CompletableFuture instantiates the self version such as new CompletableFuture.... For such methods (most of the public static methods), use an alternative method to get an instance of MDCAwareCompletableFuture. An example of using an alternative could be rather than using CompletableFuture.supplyAsync(...), you can choose new MDCAwareCompletableFuture<>().completeAsync(...)
Convert the instance of CompletableFuture to MDCAwareCompletableFuture by using the method getMDCAwareCompletionStage when you get stuck with one because of say some external library which returns you an instance of CompletableFuture. Obviously, you can't retain the context within that library but this method would still retain the context after your code hits the application code.
While supplying an executor as a parameter, make sure that it is MDC Aware such as MDCAwareForkJoinPool. You could create MDCAwareThreadPoolExecutor by overriding execute method as well to serve your use case. You get the idea!
You can find a detailed explanation of all of the above here in a post about the same.

Resources