Following up on How to use MDC with thread pools? how can one use MDC with a ForkJoinPool? Specifically, I how can one wrap a ForkJoinTask so MDC values are set before executing a task?
Following up on How to use MDC with thread pools? how can one use MDC with a ForkJoinPool? Specifically, I how can one wrap a ForkJoinTask so MDC values are set before executing a task?
Here is some additional information to go along with @Gili's answer.
Test that shows that the solution works (note that there will be lines without the Context, but at least they won't be the WRONG context, which is what was happening with a normal ForkJoinPool).
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.startsWith;
import static org.junit.Assert.assertThat;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;
import org.junit.Test;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import ch.qos.logback.classic.LoggerContext;
import ch.qos.logback.classic.spi.ILoggingEvent;
import ch.qos.logback.core.OutputStreamAppender;
public class MDCForkJoinPoolTest {
private static final Logger log = (Logger) LoggerFactory.getLogger("mdc-test");
// you can demonstrate the problem I'm trying to fix by changing the below to a normal ForkJoinPool and then running the test
private ForkJoinPool threads = new MDCForkJoinPool(16);
private Semaphore threadsRunning = new Semaphore(-99);
private ByteArrayOutputStream bio = new ByteArrayOutputStream();
@Test
public void shouldCopyManagedDiagnosticContextWhenUsingForkJoinPool() throws Exception {
for (int i = 0 ; i < 100; i++) {
Thread t = new Thread(simulatedRequest(), "MDC-Test-"+i);
t.setDaemon(true);
t.start();
}
// set up the appender to grab the output
LoggerContext lc = (LoggerContext) LoggerFactory.getILoggerFactory();
OutputStreamAppender<ILoggingEvent> appender = new OutputStreamAppender<>();
LogbackEncoder encoder = new LogbackEncoder();
encoder.setPattern("%X{mdc_val:-}=%m%n");
encoder.setContext(lc);
encoder.start();
appender.setEncoder(encoder);
appender.setImmediateFlush(true);
appender.setContext(lc);
appender.setOutputStream(bio);
appender.start();
log.addAppender(appender);
log.setAdditive(false);
log.setLevel(Level.INFO);
assertThat("timed out waiting for threads to complete.", threadsRunning.tryAcquire(300, TimeUnit.SECONDS), is(true));
Set<String> ids = new HashSet<>();
try (BufferedReader r = new BufferedReader(new InputStreamReader(new ByteArrayInputStream(bio.toByteArray()), Charset.forName("utf8")))) {
r.lines().forEach(line->{
System.out.println(line);
String[] vals = line.split("=");
if (!vals[0].isEmpty()) {
ids.add(vals[0]);
assertThat(vals[1], startsWith(vals[0]));
}
});
}
assertThat(ids.size(), is(100));
}
private Runnable simulatedRequest() {
return () -> {
String id = UUID.randomUUID().toString();
MDC.put("mdc_val", id);
Map<String, String> context = MDC.getCopyOfContextMap();
threads.submit(()->{
MDC.setContextMap(context);
IntStream.range(0, 100).parallel().forEach((i)->{
log.info("{} - {}", id, i);
});
}).join();
threadsRunning.release();
};
}
}
Also, here are the additional methods that should be overridden in the original answer.
@Override
public <T> ForkJoinTask<T> submit(ForkJoinTask<T> task) {
return super.submit(wrap(task, MDC.getCopyOfContextMap()));
}
@Override
public <T> ForkJoinTask<T> submit(Callable<T> task) {
return super.submit(wrap(task, MDC.getCopyOfContextMap()));
}
@Override
public <T> ForkJoinTask<T> submit(Runnable task, T result) {
return super.submit(wrap(task, MDC.getCopyOfContextMap()), result);
}
@Override
public ForkJoinTask<?> submit(Runnable task) {
return super.submit(wrap(task, MDC.getCopyOfContextMap()));
}
private <T> Callable<T> wrap(Callable<T> task, Map<String, String> newContext)
{
return () ->
{
Map<String, String> oldContext = beforeExecution(newContext);
try
{
return task.call();
}
finally
{
afterExecution(oldContext);
}
};
}