全链路追踪必备组件之 TransmittableThreadLocal 详解

作者 | 姜日游

杏仁 Java 工程师。今日持续在线的程序玩家。

我们都知道 ThreadLocal 作为一种多线程处理手段，将数据限制在当前线程中，避免多线程情况下出现错误。

一般的使用场景大多会是服务上下文、分布式日志跟踪。

但是在业务代码中，为了提高响应速度，将多个复杂、长时间的计算或调用过程异步进行，让主线程可以先进行其他操作。像我们项目中最常用的就是 CompletableFuture 了，默认会使用预设的 ForkJoin ThreadPool 执行。

这也就引入了一个问题，如果保证 ThreadLocal 的信息能够传递异步线程？通过 ThreadLocal？通过线程池？通过 Runnable 或者 Callable？

有些场景丢了就丢了，比如目前我们的服务上下文传递，一般都没有很严谨的处理 ......

但是，如果是分布式追踪的场景，丢了就要累惨了。

注：以下代码仅保留关键代码，其余无关紧要则忽略

InheritableThreadLocal

InheritableThreadLocal 是 JDK 本身自带的一种线程传递解决方案。顾名思义，由当前线程创建的线程，将会继承当前线程里 ThreadLocal 保存的值。

其本质上是 ThreadLocal 的一个子类，通过覆写父类中创建初始化的相关方法来实现的。我们知道，ThreadLocal 实际上是 Thread 中保存的一个 ThreadLocalMap 类型的属性搭配使用才能让广大 Javaer 直呼真香的，所以 InheritableThreadLocal 也是如此。

public class Thread implements Runnable {
    // 如果单纯使用 ThreadLocal，则 Thread 使用该属性值保存 ThreadLocalMap
    ThreadLocal.ThreadLocalMap threadLocals = null;
        // 否则使用该属性值
    ThreadLocal.ThreadLocalMap inheritableThreadLocals = null;
  
    private void init(ThreadGroup g, Runnable target, String name, long stackSize, AccessControlContext acc) {
          Thread parent = currentThread();

          if (parent.inheritableThreadLocals != null)
              this.inheritableThreadLocals =
                  ThreadLocal.createInheritedMap(parent.inheritableThreadLocals);
      }
}

init 方法作为 Thread 初始化的核心方法，相关 ThreadLocal 代码已经全部摘出。如我们所见，仅仅就只是这一点改动。在创建线程时，如果当前线程的 inheritableThreadLocals 不为空，则根据它创建出新的 InheritableThreadLocals 保存到新线程中。

Ps : ThreadLocal 作为老牌选手，默认都是使用时，直接初始化 Thread 的 threadLocals 属性。

只有像是 InheritableThreadLocal 这样的后辈，需要特殊处理一下。

public class InheritableThreadLocal extends ThreadLocal {
    
    protected T childValue(T parentValue) {
        return parentValue;
    }
  
    ThreadLocalMap getMap(Thread t) {
       return t.inheritableThreadLocals;
    }

   // Thread 中 ThreadLocalMap 不存在时的初始化动作，需要改为初始化 inheritableThreadLocals
    void createMap(Thread t, T firstValue) {
        t.inheritableThreadLocals = new ThreadLocalMap(this, firstValue);
    }
}

因此，原先 ThreadLocal 会从 Thread 的 threadLocals 获取 Map，那么 InheritableThreadLocal 就要从 inheritableThreadLocals 拿了。childValue 方法用作从父线程中获取值，可以看到，这边是直接返回的，如果是复杂对象，就直接传引用了。当然，继承覆写该方法，可以实现浅拷贝、深拷贝等等方式。

缺点

这样的方式解决了创建线程时的 ThreadLocal 传值的问题，但不可能一直创建新的线程，那实在耗费资源。因此通用做法是线程复用，比如线程池呗。但是，递交异步任务是相应的 ThreadLocal 的值就无法传递过去了。

我们希望的是，异步线程执行任务的所使用的 ThreadLocal 值，是将任务提交给线程时主线程持有的。即从任务创建时传递到任务执行时。

想想，如果我们在创建异步任务时，在任务代码外获取当前线程的值临时保存，再传递给执行线程，在真正的任务执行前保存到当前线程即可。对，确实可以，但是麻烦不？每个创建异步任务的地方都要写。

那就把它封装到递交任务的方法中。

RunnableWrapper & CallableWrapper

假设按照服务上下文的场景举例，目前项目中的执行异步操作的方案是定义一个 AsyncExecutor，并声明执行 Supplier 返回 CompletableFuture 的方法。

既然这样就可以对方法做一些改造，保证上下文的传递。

private static ThreadLocal contextHolder = new ThreadLocal<>();

public static  CompletableFuture invokeToCompletableFuture(Supplier supplier, String errorMessage) {
    // 第一步
    String context = contextHolder.get();
    Supplier newSupplier = () -> {
         // 第二步
        String origin = contextHolder.get();
        try {
            contextHolder.set(context);
            // 第三步
            return supplier.get();
        } finally {
            // 第四步
            contextHolder.set(origin);
            log.info(origin);
        }
    };
    return CompletableFuture.supplyAsync(newSupplier).exceptionally(e -> {
        throw new ServerErrorException(errorMessage, e);
    });
}
// test code
public static void main(String[] args) throws ExecutionException, InterruptedException {
    contextHolder.set("main");
    log.info(contextHolder.get());
    CompletableFuture context = invokeToCompletableFuture(() -> test.contextHolder.get(), "error");
    log.info(context.get());
}

总得来说，就是在将异步任务派发给线程池时，对其做一下上下文传递的处理。

第一步：主线程获取上下文，传递给任务暂存。

1 之后的操作都将是异步执行线程操作的。

第二步：异步执行线程将原有上下文取出，暂时保存。并将主线程传递过来的上下文设置。

第三步：执行异步任务

第四步：将原有上下文设置回去。

可以看到一般并不会在异步线程执行完任务之后直接进行 remove。而是一开始取出原上下文（可能为 NULL，也可能是线程创建时 InheritableThreadLocal 继承过来的值。当然后续也会被清除的），并在任务执行结束重新放回。这样的方式可以说是异步 ThreadLocal 传递的标准范式（大佬说的）。

这样子既起到了显式清除主线程带来的上下文，也避免了如果线程池的拒绝策略为 CallerRunsPolicy，后续处理时上下文丢失的问题。

Supplier 不算是典型例子，更为典型的应该是 Runnable 和 Callable。不过举一推三，都是修饰一下，再丢给线程池。

public final class DelegatingContextRunnable implements Runnable {

    private final Runnable delegate;

    private final Optional delegateContext;

    public DelegatingContextRunnable(Runnable delegate,
                                       Optional context) {
        assert delegate != null;
        assert context != null;

        this.delegate = delegate;
        this.delegateContext = context;
    }

    public DelegatingContextRunnable(Runnable delegate) {
        // 修饰原有的任务，并保存当前线程的值
        this(delegate, ContextHolder.get());
    }

    public void run() {
        Optional originalContext = ContextHolder.get();

        try {
            ContextHolder.set(delegateContext);
            delegate.run();
        } finally {
            ContextHolder.set(originalContext);
        }
    }
}

public final void execute(Runnable task) {
  // 递交给真正的执行线程池前，对任务进行修饰
  executor.execute(wrap(task));
}

protected final Runnable wrap(Runnable task) {
  return new DelegatingContextRunnable(task);
}

后续，使用线程池执行异步任务的时候，事先对任务进行封装代理即可。

不过，还是比较麻烦。自定义的线程池，需要显式处理任务。而且更严谨的做法，不同业务场景之间的线程池应该是隔离的，以免受到影响，就比如 Hystrix 的线程池。

每一个线程池都要处理就麻烦了。所以换个思路，代理线程池。

DelegaingExecutor

这个就不多说了，实际很简单，就照搬我们上下文相关类库。

public class DelegatingContextExecutor implements Executor  {

    private final Executor delegate;


    public DelegatingContextExecutor(Executor delegateExecutor) {
        this.delegate = delegateExecutor;
    }

    public final void execute(Runnable task) {
        delegate.execute(wrap(task));
    }

    protected final Runnable wrap(Runnable task) {
        return new DelegatingContextRunnable(task);
    }

    protected final Executor getDelegateExecutor() {
        return delegate;
    }
}
// 自定义的线程池，用于执行项目中的异步任务
public Executor queryExecutor() {
    ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor();
    // 封装服务上下文的线程池修饰
    return new DelegatingContextExecutorService(threadPoolExecutor);
}

问题似乎都解决了，那还有什么？

对，适用场景不够通用。上面的做法只针对于指定的 ThreadLocal，其他场景例如链路追踪、应用容器或上层框架跨应用代码给下层 SDK 传递信息（像是契约包 Feign 的执行线程）。

那么 TransmittableThreadLocal 就是为了解决通用化场景而设计的。

TransmittableThreadLocal

作为一个核心代码不超过一千行的工具框架，实际使用和架构设计都十分简单。

其使用方法本质上与上述提到的 CallableWrapper 和 DelegatingExecutor 是一样的，并且为了方便使用，对外提供了静态工厂方法或工具类。

public final void execute(Runnable task) {
  executor.execute(TtlCallable.get(task));
}
// 或者
public Executor queryExecutor() {
    ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor();
    // 封装服务上下文的线程池修饰
    return TtlExecutors.getTtlExecutorService(threadPoolExecutor);
}

当然，前提是 ThreadLocal 必须使用 TransmittableThreadLocal。至于为什么，我们源码分析时再细细说来。

先看看核心实现类的结构，以 Callable 和 ExecutorService 为例。

整体主要是三个部分：任务（TtlCallable）、线程池（ExecutorServiceTtlWrapper）、ThreadLocal（TransmittableThreadLocal）。其实对应上述讲到的 CallableWrapper、DelegatingExecutor、InheritableThreadLocal。

但是无论是任务和线程池，本身还是依赖于 TransmittableThreadLocal 对于存储值的管理。

用官方的时序图直观展示一下，框架是如何起作用的：

可以看到，从第三步创建完任务，第四步修饰完任务，后续大部分过程都依赖于 TransmittableThreadLocal 或 TransmittableThreadLocal 中声明的静态工具类 Transmitter。Transmitter 主要负责 ThreadLocal 的管理和值的传递。

首先看看 TtlCallable。

TtlCallable

该类实际上是 JDK Callable 的一个修饰。类比于，上文讲到的 RunnableWrapper，只是为了临时保存父线程 ThreadLocal 的值，以便在执行任务之前，赋值到子线程中。

因此，TtlCallable 和 TtlExecutorService 都实现了 TtlWrapper 接口。也许你以为，该接口是实现修饰的语义，但是它只提供了一个方法，表达了拆修饰的语义：

public interface TtlWrapper extends TtlEnhanced {
    @NonNull
    T unwrap();
}

毕竟核心是修饰，所以该类主要为了提供修饰的核心抽象，便于框架对其进行判断和管理。

该方法语义要求，必须返回修饰的源对象或下层对象（毕竟可能修饰了很多层），因此也是空值安全的。null 进来，null 出去。

public final class TtlCallable implements Callable, TtlWrapper>, TtlEnhanced, TtlAttachments {
    // 保存父线程的 ThreadLocal 快照
    private final AtomicReference