MDC+TransmittableThreadLocal实现链路追踪Demo

前言

我们知道MDC是slf4j提供给我们的扩展点，可以进行一些常用操作。比如做日志链路跟踪时，动态配置用户自定义的一些信息 但是它有一个痛点，如下图所示：简单翻译下就是子线程不能获取父线程中的数据

而【TransmittableThreadLocal】正是为了解决这样类似的通用化场景设计的

MDC-Demo

下面我们通过一个简单的demo来验证下

新建maven项目

依赖如下

    4.0.0

    org.example
    log-mdc-demo
    1.0-SNAPSHOT

    
        8
        8
    

    
        spring-boot-starter-parent
        org.springframework.boot
        2.3.8.RELEASE
    
    
        
            org.springframework.boot
            spring-boot-starter-web
        
    

application.yml如下

注意[%X{UID}] [%X{traceId}]是关键

server:
    port: 8081
spring:
    application:
        name: log-mdc


logging:
    file:
        name: log-mdc.log
        path: /logs
    pattern:
        console: '%d{yyyy-MM-dd HH:mm:ss} %-5level [%X{UID}] [%X{traceId}] %msg%n'
        file: '%d{yyyy-MM-dd HH:mm:ss} %-5level [%X{UID}] [%X{traceId}] %msg%n'

建立启动类

package com.jarvan.demo;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

import java.util.UUID;

/**
 * 描述：
 *
 * @author 朱佳文
 * @email [email protected]
 * @date 2021/3/28 2:54 下午
 * @company 数海掌讯
 */
@SpringBootApplication
public class LogMDCApp {
    private static final Logger logger = LoggerFactory.getLogger(LogMDCApp.class);

    public static void main(String[] args) {
        SpringApplication.run(LogMDCApp.class, args);
        MDC.put("traceId", UUID.randomUUID().toString());
        MDC.put("UID", "用户id");
        logger.info("===========log-mdc测试===========");
    }
}

结果

启动项目，日志输出用户id和我们的链路id

测试子线程

main方法加入下面代码

new Thread(() -> {
            String traceId = MDC.get("traceId");
            logger.info("从MDC中获取链路id为{}", traceId);
            logger.info("test runnable traceId");
        }).start();

结果如下，可以看到，子线程是无法获取父线程中的副本的

TransmittableThreadLocal

引入jar包

            com.alibaba
            transmittable-thread-local
            2.11.5
            compile
        

重写{LogbackMDCAdapter}类

package org.slf4j;

import ch.qos.logback.classic.util.LogbackMDCAdapter;
import com.alibaba.ttl.TransmittableThreadLocal;
import org.slf4j.spi.MDCAdapter;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 描述：
 * 重写{@link LogbackMDCAdapter}类，搭配TransmittableThreadLocal实现父子线程之间的数据传递
 * 内容直接从{@link LogbackMDCAdapter}类中copy。把copyOnThreadLocal实例化对象更改为TransmittableThreadLocal即可
 *
 * @author 朱佳文
 * @email [email protected]
 * @date 2021/3/28 4:09 下午
 * @company 数海掌讯
 */
public class TtlMDCAdapter implements MDCAdapter {
    final ThreadLocal> copyOnThreadLocal = new InheritableThreadLocal();
    private static final int WRITE_OPERATION = 1;
    private static final int MAP_COPY_OPERATION = 2;
    final ThreadLocal lastOperation = new ThreadLocal();
    private static TtlMDCAdapter mtcMDCAdapter;
    static {
        mtcMDCAdapter = new TtlMDCAdapter();
        MDC.mdcAdapter = mtcMDCAdapter;
    }

    public TtlMDCAdapter() {
    }
    public static MDCAdapter getInstance() {
        return mtcMDCAdapter;
    }

    private Integer getAndSetLastOperation(int op) {
        Integer lastOp = (Integer)this.lastOperation.get();
        this.lastOperation.set(op);
        return lastOp;
    }

    private boolean wasLastOpReadOrNull(Integer lastOp) {
        return lastOp == null || lastOp == MAP_COPY_OPERATION;
    }

    private Map duplicateAndInsertNewMap(Map oldMap) {
        Map newMap = Collections.synchronizedMap(new HashMap());
        if (oldMap != null) {
            synchronized(oldMap) {
                newMap.putAll(oldMap);
            }
        }

        this.copyOnThreadLocal.set(newMap);
        return newMap;
    }

    public void put(String key, String val) throws IllegalArgumentException {
        if (key == null) {
            throw new IllegalArgumentException("key cannot be null");
        } else {
            Map oldMap = (Map)this.copyOnThreadLocal.get();
            Integer lastOp = this.getAndSetLastOperation(WRITE_OPERATION);
            if (!this.wasLastOpReadOrNull(lastOp) && oldMap != null) {
                oldMap.put(key, val);
            } else {
                Map newMap = this.duplicateAndInsertNewMap(oldMap);
                newMap.put(key, val);
            }

        }
    }

    public void remove(String key) {
        if (key != null) {
            Map oldMap = (Map)this.copyOnThreadLocal.get();
            if (oldMap != null) {
                Integer lastOp = this.getAndSetLastOperation(WRITE_OPERATION);
                if (this.wasLastOpReadOrNull(lastOp)) {
                    Map newMap = this.duplicateAndInsertNewMap(oldMap);
                    newMap.remove(key);
                } else {
                    oldMap.remove(key);
                }

            }
        }
    }

    public void clear() {
        this.lastOperation.set(WRITE_OPERATION);
        this.copyOnThreadLocal.remove();
    }

    public String get(String key) {
        Map map = (Map)this.copyOnThreadLocal.get();
        return map != null && key != null ? (String)map.get(key) : null;
    }

    public Map getPropertyMap() {
        this.lastOperation.set(MAP_COPY_OPERATION);
        return (Map)this.copyOnThreadLocal.get();
    }

    public Set getKeys() {
        Map map = this.getPropertyMap();
        return map != null ? map.keySet() : null;
    }

    public Map getCopyOfContextMap() {
        Map hashMap = copyOnThreadLocal.get();
        return hashMap == null ? null : new HashMap(hashMap);
    }

    public void setContextMap(Map contextMap) {
        this.lastOperation.set(WRITE_OPERATION);
        Map newMap = Collections.synchronizedMap(new HashMap());
        newMap.putAll(contextMap);
        this.copyOnThreadLocal.set(newMap);
    }
}

更改spring.factories

org.springframework.context.ApplicationContextInitializer=\
com.jarvan.demo.config.TtlMDCAdapterInitializer

TtlMDCAdapterInitializer类如下

package com.jarvan.demo.config;

import org.slf4j.TtlMDCAdapter;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.ConfigurableApplicationContext;

/**
 * 描述：
 *
 * @author 朱佳文
 * @email [email protected]
 * @date 2021/3/28 3:59 下午
 * @company 数海掌讯
 */
public class TtlMDCAdapterInitializer implements ApplicationContextInitializer {
    @Override
    public void initialize(ConfigurableApplicationContext configurableApplicationContext) {
        TtlMDCAdapter.getInstance();
    }
}

包装线程池

public static void test() {
        ExecutorService executorService = Executors.newCachedThreadPool();
        TransmittableThreadLocal context = new TransmittableThreadLocal<>();

        context.set("test code");

        Runnable ttlRunnable = TtlRunnable.get(() -> System.out.println(context.get()));
        executorService.submit(ttlRunnable);
    }

    public static void test1() throws Exception {
        ExecutorService executorService = Executors.newCachedThreadPool();
        executorService = TtlExecutors.getTtlExecutorService(executorService);


        TransmittableThreadLocal context = new TransmittableThreadLocal<>();
        context.set("test value");

        Runnable task = () -> System.out.println("Runnable" + context.get());
        Callable call = () -> "Callable" + context.get();
        executorService.submit(task);
        executorService.submit(call);


        // --------------------------------------
        System.out.println(call.call());
        System.out.println(context.get());
    }

重新启动项目

可以看到获取到了父线程中的数据

在异步场景下有效吗

有兴趣的可以文末拉下GitHub代码测试下~

疑问点

实际使用中发现 InheritableThreadLocal 也可以实现此场景

那么为什么使用 TransmittableThreadLocal 呢？mark一下

源码

InheritableThreadLocal

先看InheritableThreadLocal ，InheritableThreadLocal只有三行代码，实际起作用的 是在thRead中的init方法中

可以看到InheritableThreadLocal是在Thread中把变量copy一份给自己了

根据它的实现，我们也可以看到它的缺点，就是 Thread的 init 方法是在线程构造方法中 copy的，也就是在线程复用的线程池中是没有办法使用的。这也解决了上面为什么要用TransmittableThreadLocal的问题

TransmittableThreadLocal

ttl使用了装饰器模式， InheritableThreadLocal 只是在线程创建的时候复制一份父线程数据，而ttl在 Thread 的 run 方法之前把父线程的数据进行了复制。核心代码如下

官方说法：

从capture()重放捕获的TransmittableThreadLocal和已注册的ThreadLocal值，并在重放之前在当前线程中返回备份的TransmittableThreadLocal值

代码下载

代码已放到GitHub，有兴趣的小伙伴可以看下~

https://github.com/JarvanBest/log-mdc-demo.git

参考：http://logback.qos.ch/manual/mdc.html