Регистратор MDC с CompletableFuture

Создать метод-оболочку

static CompletableFuture<Void> myMethod(Runnable runnable) {
    Map<String, String> previous = MDC.getCopyOfContextMap();
    return CompletableFuture.runAsync(() -> {
        MDC.setContextMap(previous);
        try {
            runnable.run();
        } finally {
            MDC.clear();
        }
    });
}

и используйте его вместо CompletableFuture.runAsync.

Моей темой решения было бы (это будет работать с JDK 9+, так как с этой версии выставлено несколько переопределяемых методов)

Ознакомьте всю экосистему с MDC

А для этого нам нужно рассмотреть следующие сценарии:

• Когда все мы получим новые экземпляры CompletableFuture из этого класса? → Нам нужно скорее вернуть ту же версию, поддерживающую MDC.
• Когда все мы получим новые экземпляры CompletableFuture из-за пределов этого класса? → Нам нужно вернуть версию того же самого с поддержкой MDC.
• Какой исполнитель используется в классе CompletableFuture? → При любых обстоятельствах мы должны убедиться, что все исполнители поддерживают MDC.

Для этого давайте создадим класс версии CompletableFuture с поддержкой MDC, расширив его. Моя версия этого будет выглядеть ниже

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

Класс MDCAwareForkJoinPool будет выглядеть так (для простоты пропущены методы с параметрами ForkJoinTask)

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

Утилиты для переноса будут такими, как

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 {
            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);
    }
}

Ниже приведены некоторые рекомендации по использованию:

• Используйте класс MDCAwareCompletableFuture, а не класс CompletableFuture.
• Несколько методов в классе CompletableFuture создают экземпляр собственной версии, такой как new CompletableFuture.... Для таких методов (большинство общедоступных статических методов) используйте альтернативный метод для получения экземпляра MDCAwareCompletableFuture. Примером использования альтернативы может быть вместо использования CompletableFuture.supplyAsync(...), вы можете выбрать new MDCAwareCompletableFuture<>().completeAsync(...)
• Преобразуйте экземпляр CompletableFuture в MDCAwareCompletableFuture с помощью метода getMDCAwareCompletionStage, когда вы застряли с одним, скажем, из-за какой-то внешней библиотеки, которая возвращает вам экземпляр CompletableFuture. Очевидно, что вы не можете сохранить контекст в этой библиотеке, но этот метод все равно сохранит контекст после того, как ваш код попадет в код приложения.
• При указании исполнителя в качестве параметра убедитесь, что он поддерживает MDC, например MDCAwareForkJoinPool. Вы также можете создать MDCAwareThreadPoolExecutor, переопределив метод execute для вашего варианта использования. Вы поняли идею!

Вы можете найти подробное объяснение всего вышеперечисленного здесь, в post о том же.

При этом ваш код может выглядеть так

new MDCAwareCompletableFuture<>().completeAsync(() -> {
            getAcountDetails(user);
            return null;
        });