如何在 CompletableFuture 中保留 slf4j MDC 日志上下文？

Question

当执行 async CompletableFuture ，父org.slf4j.MDC上下文以及org.slf4j.MDC上下文都丢失了。

这很糟糕，因为我使用某种“鱼标记”来跟踪来自多个日志文件中一个请求的日志。

MDC.put("fishid", randomId())

问题：一般而言，我如何在CompletableFutures的任务期间保留该 id？

List<CompletableFuture<UpdateHotelAllotmentsRsp>> futures =
    tasks.stream()
        .map(task -> CompletableFuture.supplyAsync(
            () -> businesslogic(task))
        .collect(Collectors.toList());

List results = futures.stream()
    .map(CompletableFuture::join)
    .collect(Collectors.toList());

public void businesslogic(Task task) {
       LOGGER.info("mdc fishtag context is lost here");
}

Answer 1

我解决这个问题的最易读的方法如下 -

---------------线程工具类--------------------

public static Runnable withMdc(Runnable runnable) {
    Map<String, String> mdc = MDC.getCopyOfContextMap();
    return () -> {
        MDC.setContextMap(mdc);
        runnable.run();
    };
}

public static <U> Supplier<U> withMdc(Supplier<U> supplier) {
    Map<String, String> mdc = MDC.getCopyOfContextMap();
    return (Supplier) () -> {
        MDC.setContextMap(mdc);
        return supplier.get();
    };
}

---------------用法--------------

CompletableFuture.supplyAsync(withMdc(() -> someSupplier()))
                 .thenRunAsync(withMdc(() -> someRunnable())
                 ....

ThreadUtils 中的 WithMdc 必须重载以包含 CompletableFuture 接受的其他功能接口

请注意 withMdc() 方法是静态导入的以提高可读性。

Answer 2

最后，我创建了一个保留MDC的Supplier包装器。 如果有人有更好的想法，请随时发表评论。

public static <U> CompletableFuture<U> supplyAsync(Supplier<U> supplier, Executor executor) {
    return CompletableFuture.supplyAsync(new SupplierMDC(supplier), executor);
}

private static class SupplierMDC<T> implements Supplier<T> {
    private final Supplier<T> delegate;
    private final Map<String, String> mdc;

    public SupplierMDC(Supplier<T> delegate) {
        this.delegate = delegate;
        this.mdc = MDC.getCopyOfContextMap();
    }

    @Override
    public T get() {
        MDC.setContextMap(mdc);
        return delegate.get();
    }
}

Answer 3

我的解决方案主题是（它可以与 JDK 9+ 一起使用，因为自该版本以来公开了几个可覆盖的方法）

让整个生态系统了解 MDC

为此，我们需要解决以下场景：

• 我们什么时候才能从这个类中获得 CompletableFuture 的新实例？ → 我们需要返回相同的 MDC 感知版本。
• 我们什么时候才能从这个类之外获得 CompletableFuture 的新实例？ → 我们需要返回相同的 MDC 感知版本。
• 在 CompletableFuture 类中使用哪个执行器？ → 在任何情况下，我们都需要确保所有执行者都知道 MDC

为此，让我们通过扩展它来创建一个 MDC 感知版本的CompletableFuture类。 我的版本如下所示

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

以下是一些使用指南：

• 使用类MDCAwareCompletableFuture而不是类CompletableFuture 。
• CompletableFuture类中的几个方法实例化了自我版本，例如new CompletableFuture... 。 对于此类方法（大多数公共静态方法），请使用替代方法来获取MDCAwareCompletableFuture的实例。 使用替代方案的一个例子可能是而不是使用CompletableFuture.supplyAsync(...) ，您可以选择new MDCAwareCompletableFuture<>().completeAsync(...)
• 使用MDCAwareCompletableFuture方法将CompletableFuture的实例转换为getMDCAwareCompletionStage当您因为某个外部库返回CompletableFuture实例而陷入getMDCAwareCompletionStage时。 显然，您不能在该库中保留上下文，但是在您的代码命中应用程序代码后，此方法仍会保留上下文。
• 在提供 executor 作为参数时，请确保它是 MDC Aware，例如MDCAwareForkJoinPool 。 您也可以通过覆盖execute方法来创建MDCAwareThreadPoolExecutor来为您的用例提供服务。 你明白了！

有了这个，你的代码看起来像

List<CompletableFuture<UpdateHotelAllotmentsRsp>> futures =
    tasks.stream()
        new MDCAwareCompletableFuture<UpdateHotelAllotmentsRsp>().completeAsync(
            () -> businesslogic(task))
        .collect(Collectors.toList());

List results = futures.stream()
    .map(CompletableFuture::join)
    .collect(Collectors.toList());

public UpdateHotelAllotmentsRsp businesslogic(Task task) {
       LOGGER.info("mdc fishtag context is not lost here");
}

您可以在相同的帖子中找到对上述所有内容的详细说明。

Answer 4

是的，Twitter Future 正确地做到了这一点。 他们有一个 Future.scala 知道的 Local.scala 类。

该修复程序供 Java 作者修复此问题，以便您的本地状态通过所有使用 CompletableFutures 的库。 基本上，Local.scala 由 Future 使用，并且在内部使用 ThreadLocal 直到 .thenApply 或 .thenAccept 它将捕获状态并在需要时将其转移到下一个状态。 这适用于所有第三方库，零 3rd 方库更改。

这里还有更多，但戳 Java 作者来修复他们的东西...... http://mail.openjdk.java.net/pipermail/core-libs-dev/2017-May/047867.html

在那之前，MDC 永远不会通过 3rd 方库工作。

我的 SO 帖子CompletableFuture 有相应的本地上下文吗？