线程池如何传递线程上下文信息

业务开发中，一般都会使用ThreadLocal保存一些上下文信息，但是在线程池中执行对应逻辑时，由于是不同线程所以无法获取之前线程的上下文信息。

线程池的线程上下文传递，实现方案就是在提交任务时记录当前线程上下文信息，在线程池中线程执行用户任务前将之前保存的上下文塞到当前线程的上下文中，在执行用户任务之后移除该上下文即可。简单来说就是，外部线程提交任务时要记录上下文信息，内部线程执行任务时获取之前记录的上下文信息设置到当前线程上下文中。

实现线程上下文传递的2种方式：

一种是在用户任务中直接进行手动获取/设置上下文逻辑。
另一种是实现一个自定义的线程池，在提交任务时对任务进行包装并保存上下文信息，然后任务执行前设置上下文信息。

两种实现方式的代码如下：

private static ThreadLocal<String> CONTEXT = new ThreadLocal<>();
private static ExecutorService executor = new ThreadPoolExecutor(1, 1,
        60, TimeUnit.SECONDS, new ArrayBlockingQueue<>(512));

private static ExecutorService executorWrap = new ThreadPoolExecutorWrap(1, 1,
        60, TimeUnit.SECONDS, new ArrayBlockingQueue<>(512));

public static void main(String[] args) {
    CONTEXT.set("main context");

    // 方式1：在用户任务中直接进行手动获取/设置上下文逻辑
    executor.submit(new RunnableWrap(() -> System.out.println("hello world: " + CONTEXT.get())));

    // 方式2：自定义线程池，封装成支持保存/设置上下文的任务
    executorWrap.submit(() -> System.out.println("hello world: " + CONTEXT.get()));
}

static class ThreadPoolExecutorWrap extends ThreadPoolExecutor {
    public ThreadPoolExecutorWrap(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    @Override
    public Future<?> submit(Runnable task) {
        if (task == null) {
            throw new NullPointerException();
        }
        RunnableFuture<Void> ftask = newTaskFor(new RunnableWrap(task), null);
        execute(ftask);
        return ftask;
    }
}

static class RunnableWrap implements Runnable {
    private String contextValue;
    private Runnable task;

    public RunnableWrap(Runnable task) {
        this.contextValue = CONTEXT.get();
        this.task = task;
    }

    @Override
    public void run() {
        try {
            CONTEXT.set(contextValue);
            // 用户任务逻辑
            task.run();
        } finally {
            CONTEXT.remove();
        }
    }
}

关于线程间上下文传递，阿里给出了一个解决方案：TTL（transmittable-thread-local）是一个线程间传递ThreadLocal，异步执行时上下文传递的解决方案。整个库的核心是构建在TransmittableThreadLocal类（继承并加强InheritableThreadLocal类）之上，同时包含线程池修饰（ExecutorService/ForkJoinPool/TimerTask）以及Java Agent支持，代码小于1k行，短小精悍。

我们都知道，JDK的InheritableThreadLocal类可以完成父线程到子线程的值传递。但对于使用线程池等会池化复用线程的组件的情况，线程由线程池创建好，并且线程是池化起来反复使用的；这时父子线程关系的ThreadLocal值传递已经没有意义，应用需要的实际上是把 任务提交给线程池时的ThreadLocal值传递到 任务执行时。原理是使用TtlRunnable/Ttlcallable包装了Runnable/Callable类：

在TtlRunnable/Ttlcallable初始化时capture TransmittableThreadLocal变量
在run方法调用runnable.run()前进行replay，设置到当前线程ThreadLocal
在run方法调用runnable.run()后进行restore，上下文还原，也就是replay的反向操作

注意，步骤1和步骤2/3不是在同一个线程中执行的，这个流程和本文最初说的实现方案是一致的。

TTL的示例代码如下：

void testTtlInheritableThreadLocal() throws InterruptedException {
    ExecutorService executor = Executors.newFixedThreadPool(1);
    executor.submit(() -> {}); // 先进行工作线程创建

    // 使用TTL
    final TransmittableThreadLocal<String> parent = new TransmittableThreadLocal<>();
    parent.set("value-set-in-parent");
    // 将Runnable通过TtlRunnable包装下
    executor.submit(TtlRunnable.get(() -> System.out.println(Thread.currentThread().getName() + ": " + parent.get())));
}
// 输出结果：pool-1-thread-1: value-set-in-parent

TTL实现原理和文章开头说的实现线程上下文传递大致一致，感兴趣的小伙伴可以直接看下TTL源码（https://github.com/alibaba/transmittable-thread-local），这里不再赘述。

最后关于ThreadLocal再提一下，我们可以重写其initialValue方法，这样可以在threadLocal.get为空时初始化一个值，使用示例如下：

ThreadLocal<String> local = new ThreadLocal<String>() {
    @Override
    protected String initialValue() {
        return "init";
    }
};

System.out.println(local.get()); // init

local.set("hello world");
System.out.println(local.get()); // hello world

local.remove();
System.out.println(local.get()); // init