MDC-线程池透传traceId

Java开发中都是使用日志门面+日志实现的方式打印日志。日志门面主要是为了给Java日志访问提供一套标准、规范的API框架，其主要意义在于提供接口，具体的实现可以交由具体的日志实现框架。使用logback在多线程环境下MDC会丢失父线程的上下文。

public class LogbackMDCAdapter implements MDCAdapter  {
  final ThreadLocal<Map<String, String>> readWriteThreadLocalMap = new ThreadLocal<Map<String, String>>();
  final ThreadLocal<Map<String, String>> readOnlyThreadLocalMap = new ThreadLocal<Map<String, String>>();
  private final ThreadLocalMapOfStacks threadLocalMapOfDeques = new ThreadLocalMapOfStacks();

  ....
}

这在 logback 使用的是 ThreadLocal ，并没有使用 InheritableThreadLocal ，这是因为使用 InheritableThreadLocal 可能导致内存泄漏: ThreadLocal 的工作原理是为每个线程维护一个独立的变量副本，当使用 InheritableThreadLocal时，子线程会继承父线程的 ThreadLocal 变量，如果线程一直不结束，或者线程池中的线程一直被复用，而没有正确清理 ThreadLocal 的值，就会导致内存泄漏。

MDC

MDC(Mapped Diagnostic Context): 诊断上下文映射。用户可以通过将关心的上下文信息写入到MDC中，在日志输出时自动输出，而不需要手动的设置。MDC核心代码如下：

public class MDC {

  public static void put(String key, String val) throws IllegalArgumentException {
      if (key == null) {
          throw new IllegalArgumentException("key parameter cannot be null");
      }
      if (getMDCAdapter() == null) {
          throw new IllegalStateException(MDC_APAPTER_CANNOT_BE_NULL_MESSAGE);
      }
      getMDCAdapter().put(key, val);
  }

  public static void clear() {
      if (getMDCAdapter() == null) {
          throw new IllegalStateException(MDC_APAPTER_CANNOT_BE_NULL_MESSAGE);
      }
      getMDCAdapter().clear();
  }

  public static Map<String, String> getCopyOfContextMap() {
      if (getMDCAdapter() == null) {
          throw new IllegalStateException(MDC_APAPTER_CANNOT_BE_NULL_MESSAGE);
      }
      return getMDCAdapter().getCopyOfContextMap();
  }

  public static void remove(String key) throws IllegalArgumentException {
      if (key == null) {
          throw new IllegalArgumentException("key parameter cannot be null");
      }

      if (getMDCAdapter() == null) {
          throw new IllegalStateException(MDC_APAPTER_CANNOT_BE_NULL_MESSAGE);
      }
      getMDCAdapter().remove(key);
  }

  public static void setContextMap(Map<String, String> contextMap) {
      if (getMDCAdapter() == null) {
          throw new IllegalStateException(MDC_APAPTER_CANNOT_BE_NULL_MESSAGE);
      }
      getMDCAdapter().setContextMap(contextMap);
  }
}

slf4j中 MDCAdapter 实现，logback并没有使用slf4j的BasicMDCAdapter而是自己实现了。

public class BasicMDCAdapter implements MDCAdapter {

  private final ThreadLocalMapOfStacks threadLocalMapOfDeques = new ThreadLocalMapOfStacks();

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

多线程中使用MDC

工具类封装

public class ThreadWrapper {

  public static Runnable runnable(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();
        }
      }
    };
  }
}

实现ThreadFactory

public class DefaultThreadFactory implements ThreadFactory {

  private static final AtomicInteger poolId = new AtomicInteger(0);
  private final AtomicInteger threadId = new AtomicInteger(0);

  private final String prefix;
  private final boolean daemon;
  private final int priority;
  protected final ThreadGroup group;

  public DefaultThreadFactory(String name) {
    this(name, false);
  }

  public DefaultThreadFactory(String name, boolean daemon) {
    this(name, daemon, Thread.NORM_PRIORITY, null);
  }

  public DefaultThreadFactory(String name, boolean daemon, ThreadGroup group) {
    this(name, daemon, Thread.NORM_PRIORITY, group);
  }

  public DefaultThreadFactory(String name, boolean daemon, int priority, ThreadGroup group) {
    this.prefix = name + "-" + poolId.incrementAndGet() + "-thread-";
    this.daemon = daemon;
    this.priority = priority;
    this.group = group;
  }

  @Override
  public Thread newThread(@Nonnull Runnable runnable) {
    // 这里特殊处理，ThreadWrapper.runnable()将父线程context传给子线程
    Thread thread = new Thread(group, ThreadWrapper.runnable(runnable),
        prefix + threadId.incrementAndGet(), 0);
    if (thread.isDaemon() != daemon) {
      thread.setDaemon(daemon);
    }

    if (thread.getPriority() != priority) {
      thread.setPriority(priority);
    }
    return thread;
  }
}

核心： Thread thread = new Thread(group, ThreadWrapper.runnable(runnable), prefix + threadId.incrementAndGet(), 0)

示例

public class DefaultExecutorPoolExample {

  private static final Logger logger = LoggerFactory.getLogger(DefaultExecutorPoolExample.class);

  public static void main(String[] args) throws ExecutionException, InterruptedException {

    MDC.put("traceId", String.valueOf(Instant.now().toEpochMilli()));
    ExecutorService executorService = new ThreadPoolExecutor(10, 10, 3, TimeUnit.SECONDS,
        new ArrayBlockingQueue<>(65535), new DefaultThreadFactory("test"));

    for (int i = 0; i < 10; i++) {
      executorService.submit(() -> {
        logger.info("sub task, thread name: {}", Thread.currentThread().getName());
      });
    }
    executorService.shutdown();

    ScheduledExecutorService scheduleService = new ScheduledThreadPoolExecutor(1,
        new DefaultThreadFactory("test"));

    scheduleService.scheduleAtFixedRate(
        () -> log.info("sub task scheduleService, Thread Nmae: {}..........................",
            Thread.currentThread().getName()), 0,
        3, TimeUnit.SECONDS);
    
    MDC.remove("traceId");
  }
}