Xudong's Blog

让MDC在各种线程间穿梭自如

Java
Word count: 1.8kReading time: 7 min
2022/12/17

在Java应用中，Logback或Log4j的MDC可以帮我们实现在记录日志时携带当前上下文之外的信息。一个典型的应用场景是，当一个HTTP请求进入Spring应用后，在拦截器中向MDC写入为当前请求生成trace id / transaction id，后续业务类中打印的log都可以携带id，方便日志的定位。

MDC的实现方式

由于我们经常使用Slf4j之类的log facade，在使用时只需要关心类似org.slf4j.MDC的类提供的接口就可以。但是在具体日志库实现时，通常会实现一个MDCAdapter接口以便门面MDC类调用。不论是Logback还是Log4j，实现MDC的方式都是通过ThreadLocal来存储我们写入MDC的数据。

以Logback为例，它的LogbackMDCAdapter类中的第一行代码就是创建一个ThreadLocal对象。

1
2
3
public class LogbackMDCAdapter implements MDCAdapter {
    final ThreadLocal<Map<String, String>> copyOnThreadLocal = new ThreadLocal();
... ...

新版本的logback使用了两个ThreadLocal成员变量 readWriteThreadLocalMap和readOnlyThreadLocalMap实现读写分离，在读取上下文时不会受到写操作的影响，去除了代码中复制map时使用的synchronized关键字，原来提高了并发访问效率。但是在LogbackMDCAdapterSimple的实现中，仍然使用了单个ThreadLocal成员变量来存储上下文Map。

奇怪的是，Slf4j提供的一个MDCAdapter基础实现类org.slf4j.helpers.BasicMDCAdapter中使用了一个InheritableThreadLocal解决子线程中MDC传递的问题。这个类的注释中也提到了，代码其实是参考LogbackMDCAdapter的，现在后者更新进化后更复杂了。那么为什么Logback现在反而使用了对子线程不友好的基础ThreadLocal类呢？

1
2
3
4
5
6
7
8
9
10
11
12
public class BasicMDCAdapter implements MDCAdapter {

    private InheritableThreadLocal<Map<String, String>> inheritableThreadLocal = new InheritableThreadLocal<Map<String, String>>() {
        @Override
        protected Map<String, String> childValue(Map<String, String> parentValue) {
            if (parentValue == null) {
                return null;
            }
            return new HashMap<String, String>(parentValue);
        }
    };
... ...

结合LogbackMDCAdapter的代码和注释，google以后得知，放弃使用InheritableThreadLocal是因为logback的MDC曾因为这个类出现过内存泄露问题。并且，当子线程存在时，后续不管是父线程还是子线程调用MDC.put()方法，都会让子线程和父线程中的MDC上下文map出现差异，这可能导致调用者预期之外的结果。与其让这种预期之外的情况出现，不如让调用者自己管理线程之间的MDC传递，隔离不同线程下的MDC。

多线程下的MDC传递

跨线程传递

org.slf4j.MDC类提供了Map<String, String> getCopyOfContextMap()和setContextMap(Map<String, String> contextMap)两个方法来方便我们获取和设置整个上下文Map。所以我们可以轻易写出类似下面的代码：

1
2
3
4
5
6
7
final Map<String, String> mdcMap = MDC.getCopyOfContextMap();
new Thread(() -> {
    MDC.setContextMap(mdcMap);
    try { ... ... } finally {
        MDC.clear();  // 线程结束前清空MDC
    }
}).start();

通过lambda外的final Map变量来设置原来线程上下文的MDC非常方便。需要注意的是，一定要在线程结束前清空MDC，否则会导致内存泄露。如果是使用CompletableFuture，不要忘记在.exceptionally()的lambda里加上MDC.clear()。

此时，如果我们实现一个辅助类来帮助装饰我们的Runnable lambda，会让代码更加简洁，可维护性更强。类中的包装方法结构都是类似的，只不过包装的对象不同。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
public class MDCDecorator {
    public static <T> Callable<T> buildCallable(final Callable<T> callable) {
        final Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> previousMdcContext = MDC.getCopyOfContextMap();
            if (context == null || context.isEmpty()) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                return callable.call();
            } finally {
                if (previousMdcContext != null) {
                    // 恢复之前的MDC上下文
                    MDC.setContextMap(previousMdcContext);
                } else {
                    // 清除子线程的MDC，避免内存溢出
                    MDC.clear();
                }
            }
        };
    }

    public static Runnable buildRunnable(final Runnable runnable) {
        final Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> previousMdcContext = MDC.getCopyOfContextMap();
            if (context == null || context.isEmpty()) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                if (previousMdcContext != null) {
                    // 恢复之前的MDC上下文
                    MDC.setContextMap(previousMdcContext);
                } else {
                    // 清除子线程的MDC，避免内存溢出
                    MDC.clear();
                }
            }
        };
    }

    public static <T> Supplier<T> buildSupplier(Supplier<T> supplier) {
        final Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> previousMdcContext = MDC.getCopyOfContextMap();
            if (context == null || context.isEmpty()) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                // 调用原始Supplier
                return supplier.get();
            } finally {
                // 恢复之前的MDC上下文
                if (previousMdcContext != null) {
                    MDC.setContextMap(previousMdcContext);
                } else {
                    // 清除子线程的MDC，避免内存溢出
                    MDC.clear();
                }
            }
        };
    }
}

以后，不管是在手动创建线程时还是使用CompletableFuture，都可以很简单的带上MDC：

1
2
3
4
5
6
7
new Thread(MDCDecorator.buildRunnable(() -> {
    ... ...
})).start();

CompletableFuture.supplyAsync(MDCDecorator.buildSupplier(() -> {
    return new Object();
}));

跨线程池传递

虽然上面的MDCDecorator也可以用在向线程池提交任务时。但是，对于我们经常在Spring应用程序中使用的ThreadPoolTaskExecutor，有更优雅的方法来设置一个装饰器。首先，要实现一个org.springframework.core.task.TaskDecorator。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public class MDCTaskDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        // 这里获取的Map来自父线程
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                if (contextMap != null) {
                   // lambda是由子线程执行的，这里是在子线程内设置ContextMap
                    MDC.setContextMap(contextMap);
                }
                // 执行原本传入线程池的Runnable对象
                runnable.run();
            } finally {
                // 清除子线程的MDC 避免内存溢出
                MDC.clear();
            }
        };
    }
}

然后在初始化Executor的Bean时设置装饰器：

1
2
3
4
5
6
7
8
9
10
@Bean("taskExecutor")
public Executor taskExecutor() {
     ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
     ... ...
     // 设置异步MDC装饰器
     executor.setTaskDecorator(new MDCTaskDecorator());
     
     executor.initialize();
     return executor;
}

注意其它线程池

除了我们手动向线程池提交任务，Web服务中可能还会存在别的线程池。例如，在使用javax.websocket.server.ServerEndpoint（或jakarta.websocket.server.ServerEndpoint）注解标识的类会提供websocket接口，这个类中使用@OnMessage、@OnOpen等注解的方法会被异步地在javax.websocket实现类内部维护的线程池中被回调。

在这种情况下，我们要借助javax.websocket.Session / jakarta.websocket.Session 对象来存储MDC信息。并在onMessage、onOpen等方法内部手动把Session中存储的ContextMap拿出来设置为MDC。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
@OnOpen
public void onOpen(Session session, EndpointConfig endpointConfig) {
   // 在WebSocket链接建立时，设置新的trace-id，而不是在拦截器中
   MDC.put("traceId", TraceIDUtil.generate());
   // 保存MDC上下文信息到session中，以便后续使用
   Map<String, String> contextMap = MDC.getCopyOfContextMap();
   session.getUserProperties().put("MDC_CONTEXT_MAP", contextMap);
   ... ...
}

@OnMessage
public void onMessage(Session session, String message, @PathParam("model-type") String modelType) throws IOException {
   // 从session中恢复MDC上下文
   Map<String, String> contextMap  = (Map<String, String>) session.getUserProperties().get("MDC_CONTEXT_MAP");
   MDC.setContextMap(contextMap);
   try {
      ... ...
   } finally {
      MDC.clear();
   }
}

有些Http请求库可能也会使用内部的线程池来优化并发性能，都需要根据具体类库来定制decorator。另外，跨服务的MDC传递也还没有讨论，这些内容后面有空再写。

Reference

Java Logging with Mapped Diagnostic Context (MDC) Baeldung
实现全面生态化的全链路日志追踪系统服务插件「Logback-MDC篇」
LOGBACK-620 Performance improvement for LogbackMDCAdapter - QOS.ch JIRA

CATALOG
  1. 1. MDC的实现方式
  2. 2. 多线程下的MDC传递
    1. 2.1. 跨线程传递
    2. 2.2. 跨线程池传递
    3. 2.3. 注意其它线程池
Total : 19
2022
2021
2020
2019
2018
Hexo Github Java Linux OOP Redis Git RSS Python LaTeX Generics Note Crawler Repost Network Photograph
Tech Share