使用TransmittableThreadLocal实现异步场景日志链路追踪

• 背景
• 解决方案

背景

    在生产环境排查问题往往都是通过日志，但对于巨大的日志量，如何针对某一个操作进行一整个日志链路的追踪就显得尤为重要，在Java语言第三方的日志工具都提了日志链路追踪的方案，比如logback的MDC，MDC的使用也很简单，就是在业务的开始put一个key-value，这个key-value就能贯穿整个线程的执行流程，使用代码如下：

MDC.put("traceId", UUID.randomUUID().toString());
复制代码

    MDC虽然提供了一个现成的整个执行流程的日志追踪的方案，但是也只是一个线程，假如一个线程中又启动了另一个线程呢，这时MDC就无法完成完整的链路追踪工作了，因为MDC是基于ThreadLocal实现的，所以当一个线程中启动另一个线程的时候两个线程的TraceId就隔离开了，也就无法做到日志链路追踪。

解决方案

线程间参数传递技术选型

• InheritableThreadLocal

    InheritableThreadLocal是JDK实现的一种线程传递解决方案，由当前线程创建的线程，将会继承当前线程里ThreadLocal保存的值，但由于InheritableThreadLocal是在创建线程是解决ThreadLocal的传值问题，但是线程不可能一直创建，在工程代码中往往都是使用线程池，但是，递交异步任务使相应的ThreadLocal的值就无法传递过去了。

• TransmittableThreadLocal真正的解决方案

    TransmittableThreadLocal是阿里巴巴开源的，用于解决在使用线程池等会缓存线程的组件情况下传递ThreadLocal问题的InheritableThreadLocal扩展，具体的实现以后有机会深入研究，该解决方案常用于以下几个场景：

分布式跟踪系统
应用容器或上层框架跨应用代码给下层SDK传递信息
日志收集记录系统上下文

重写MDCAdapter

    由于MDC是基于ThreadLocal实现的，所以我们现在需要做的就是重写MDCAdapter，使系统再使用MDC时实际上是使用的我们自己实现的MDCAdapter，自定义的MDCAdapter时要注意包名应该与logback的MDCAdapter一致，因为我们要在程序启动的时候替换MDC中的MDCAdapter，MDC的MDCAdapter是包级私有，所以自定义MDCAdapter的包名一定要哥logback的MDCAdapter一致，自定义MDCAdapter代码如下：

package org.slf4j;


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;

/**
 * 重写logback的LogbackMDCAdapter，注意MDC是MDCAdapter是包级私有，所以重写的报名应该是org.slf4j
 * 用TransmittableThreadLocal替换ThreadLocal，解决多线程，异步情况x-glabal-sessionId无法传递问题
 * @author liupenghui
 * @date 2021/6/30 2:38 下午
 */
public class TtlMDCAdapter implements MDCAdapter {

    private final ThreadLocal<Map<String, String>> copyOnInheritThreadLocal = new TransmittableThreadLocal<>();

    private static final int WRITE_OPERATION = 1;
    private static final int MAP_COPY_OPERATION = 2;

    static TtlMDCAdapter mtcMDCAdapter;

    static {
        mtcMDCAdapter = new TtlMDCAdapter();
        // 替换MDC的MDCAdapter
        MDC.mdcAdapter = mtcMDCAdapter;
    }

    public static MDCAdapter getInstance() {
        return mtcMDCAdapter;
    }

    final ThreadLocal<Integer> lastOperation = new ThreadLocal<Integer>();

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

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

    private Map<String, String> duplicateAndInsertNewMap(Map<String, String> oldMap) {
        Map<String, String> newMap = Collections.synchronizedMap(new HashMap<String, String>());
        if (oldMap != null) {
            // we don't want the parent thread modifying oldMap while we are
            // iterating over it
            synchronized (oldMap) {
                newMap.putAll(oldMap);
            }
        }

        copyOnInheritThreadLocal.set(newMap);
        return newMap;
    }

    /**
     * Put a context value (the <code>val</code> parameter) as identified with the
     * <code>key</code> parameter into the current thread's context map. Note that
     * contrary to log4j, the <code>val</code> parameter can be null.
     * <p/>
     * <p/>
     * If the current thread does not have a context map it is created as a side
     * effect of this call.
     *
     * @throws IllegalArgumentException in case the "key" parameter is null
     */
    @Override
    public void put(String key, String val) throws IllegalArgumentException {
        if (key == null) {
            throw new IllegalArgumentException("key cannot be null");
        }

        Map<String, String> oldMap = copyOnInheritThreadLocal.get();
        Integer lastOp = getAndSetLastOperation(WRITE_OPERATION);

        if (wasLastOpReadOrNull(lastOp) || oldMap == null) {
            Map<String, String> newMap = duplicateAndInsertNewMap(oldMap);
            newMap.put(key, val);
        } else {
            oldMap.put(key, val);
        }
    }

    /**
     * Remove the the context identified by the <code>key</code> parameter.
     * <p/>
     */
    @Override
    public void remove(String key) {
        if (key == null) {
            return;
        }
        Map<String, String> oldMap = copyOnInheritThreadLocal.get();
        if (oldMap == null)
            return;

        Integer lastOp = getAndSetLastOperation(WRITE_OPERATION);

        if (wasLastOpReadOrNull(lastOp)) {
            Map<String, String> newMap = duplicateAndInsertNewMap(oldMap);
            newMap.remove(key);
        } else {
            oldMap.remove(key);
        }
    }

    /**
     * Clear all entries in the MDC.
     */
    @Override
    public void clear() {
        lastOperation.set(WRITE_OPERATION);
        copyOnInheritThreadLocal.remove();
    }

    /**
     * Get the context identified by the <code>key</code> parameter.
     * <p/>
     */
    @Override
    public String get(String key) {
        final Map<String, String> map = copyOnInheritThreadLocal.get();
        if ((map != null) && (key != null)) {
            return map.get(key);
        } else {
            return null;
        }
    }

    /**
     * Get the current thread's MDC as a map. This method is intended to be used
     * internally.
     */
    public Map<String, String> getPropertyMap() {
        lastOperation.set(MAP_COPY_OPERATION);
        return copyOnInheritThreadLocal.get();
    }

    /**
     * Returns the keys in the MDC as a {@link Set}. The returned value can be
     * null.
     */
    public Set<String> getKeys() {
        Map<String, String> map = getPropertyMap();

        if (map != null) {
            return map.keySet();
        } else {
            return null;
        }
    }

    /**
     * Return a copy of the current thread's context map. Returned value may be
     * null.
     */
    @Override
    public Map<String, String> getCopyOfContextMap() {
        Map<String, String> hashMap = copyOnInheritThreadLocal.get();
        if (hashMap == null) {
            return null;
        } else {
            return new HashMap<String, String>(hashMap);
        }
    }

    @Override
    public void setContextMap(Map<String, String> contextMap) {
        lastOperation.set(WRITE_OPERATION);

        Map<String, String> newMap = Collections.synchronizedMap(new HashMap<String, String>());
        newMap.putAll(contextMap);

        // the newMap replaces the old one for serialisation's sake
        copyOnInheritThreadLocal.set(newMap);
    }
}
复制代码

TtlMDCAdapter生效

• Spring程序

程序启动时设置MDCAdapter，代码如下：

TtlMDCAdapter.getInstance();
复制代码

• SpringBoot程序

实现ApplicationContextInitializer接口，重写initialize方法，代码如下：

public class TtlMDCAdapterInitializer implements ApplicationContextInitializer<ConfigurableApplicationContext> {

    @Override
    public void initialize(ConfigurableApplicationContext configurableApplicationContext) {
        // 加载自定义的MDCAdapter
        TtlMDCAdapter.getInstance();
    }
}
复制代码

TtlMDCAdapterInitializer生效的三种方式：

在META-INF目录下创建spring.factories，内容如下：

# SpringBoot程序添加自定义ApplicationContextInitializer，用以支持异步程序日志链路追踪
org.springframework.context.ApplicationContextInitializer=com.ruubypay.log.TtlMDCAdapterInitializer
复制代码

application.properties添加配置方式：

context.initializer.classes=com.ruubypay.log.TtlMDCAdapterInitializer
复制代码

程序代码addInitializers

@SpringBootApplication
public class Application {
 
	public static void main(String[] args) {
		SpringApplication springApplication = new SpringApplication(Application.class);
        springApplication.addInitializers(new Demo01ApplicationContextInitializer());
        springApplication.run(args);
	}
}
复制代码