实现ThreadLocal在异步线程（包括线程池）中传递

1.前言

我们知道通过使用ThreadLocal能够在当前线程内共享变量，该变量是和其他线程隔离的。大致的实现就是Thread类中持有了ThreadLocal类的内部类ThreadLocalMap，在调用ThreadLocal的set方法时，会先拿到当前Thread的ThreadLocalMap，再将当前ThreadLocal实例作为key设置到ThreadLocalMap中，value即为ThreadLocal的set方法设置的值，从而达到变量与当前Thread绑定的目的。

为了迎合多线程的场景，java还提供了InheritableThreadLocal类，当你在主线程内通过InheritableThreadLocal设置了值时，新建了一个子线程，在该线程内执行业务逻辑时，也可通过InheritableThreadLocal获取到主线程设置的值。原理大概就是在进行new Thread()时，会把当前线程的ThreadLocalMap拷贝到子线程中去，从而实现子线程也能获取得到主线程中设置的值。

然而目前大部分情况都不推荐直接通过new Thread()的情况执行任务，而是使用线程池来实现。由于线程池中线程是可复用的，执行每个任务不一定都会创建线程，因此上述InheritableThreadLocal在这种场景下就达不到预期的效果了。

2.问题发现

在项目中由于很多情况下都需要从请求头中获取对应的请求头信息，如手机的设备信息，为了方便会通过实现HandlerInterceptor将该拦截器注入Spring容器中，在方法执行前获取到请求头信息封装成一个对象保留到ThreadLocal中，从而在各个接口中就可通过该ThreadLocal很方便的获取到请求头信息。

public class ReqHeaderInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        CommonHeaderEntity headerEntity = CommonHeaderEntity.builder()
                .xUserAgent(request.getHeader("x-UserAgent"))
                .xNetType(request.getHeader("x-NetType"))
                .xDeviceInfo(request.getHeader("x-DeviceInfo"))
                .build();
        ThreadLocalSetting.HEADER.set(headerEntity);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        ThreadLocalSetting.HEADER.remove();
    }
}

因为在方法执行前就已经设置了变量，只有主线程执行的情况下是不用担心获取到的CommonHeaderEntity为空的。但是当引入了线程池就出问题了，获取到的CommonHeaderEntity为空，在get其中的属性时直接就抛出了NullPointException。当然有部分原因是代码缺陷，没有做判空处理，但是根本原因还是主线程的ThreadLocal变量值没有传递给子线程。

3.实现可传递的ThreadLocal

我们知道InheritableThreadLocal传递给子线程也是通过拷贝的方式将变量拷贝到子线程的InheritableThreadLocal当中，本质就是变量的拷贝，沿着这条思路走，也就可以实现在使用线程池时进行变量的传递。因为主线程与子线程间都是先通过实现一个Runnable任务开始，在创建这个Runnable实现类时，就可以将主线程的ThreadLocal变量拷贝到一个Map集合，保留到Runnable实现类的一个私有Map变量中，key为ThreadLocal实例，value为与ThreadLocal绑定的变量值。在执行Thread的run方法时，先遍历Map集合，调用每个ThreadLocal实例的set方法，set的值就是与当前ThreadLocal绑定的变量值，这样就实现了两个线程间的传递。实现如下：

public class ThreadLocalTransferDecorator {

    public static Runnable decorate(Runnable runnable) {
        return new ThreadLocalTransferRunnable(runnable);
    }

    private static class ThreadLocalTransferRunnable implements Runnable {

        private Runnable runnable;
        private Map<ThreadLocal, Object> context;
        private Thread mainThread;

        ThreadLocalTransferRunnable(Runnable runnable) {
            this.runnable = runnable;
            mainThread = Thread.currentThread();
            Map<ThreadLocal<?>, Object> tlMap = TransferableThreadLocal.getTlMap();
            context = new HashMap<>(tlMap);
        }

        private void contextTransfer() {
            context.forEach(ThreadLocal::set);
            context = null;
        }

        @Override
        public void run() {
            contextTransfer();
            try {
                runnable.run();
            } finally {
                TransferableThreadLocal.clear(mainThread);
            }
        }
    }
}

看到这你会想TransferableThreadLocal是什么时候记录下这些映射关系的？我们可以通过继承InheritableThreadLocal并重写它的initialValue、set、remove方法，在设置和删除的时候记录下的映射关系，THREAD_LOCAL_MAP存储的是与当前线程绑定的ThreadLocal集合。在上述Runnable包装类执行完成时，通过调用clear方法清理掉与当前线程绑定的所有ThreadLocal实例，更加的方便，从而避免了手动remove。

public class TransferableThreadLocal<T> extends InheritableThreadLocal<T> {

    private static final ThreadLocal<Map<ThreadLocal<?>, Object>> THREAD_LOCAL_MAP = ThreadLocal.withInitial(WeakHashMap::new);

    @Override
    protected T initialValue() {
        T value = super.initialValue();
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.put(this, value);
        return value;
    }

    @Override
    public void set(T value) {
        super.set(value);
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.put(this, value);
    }

    @Override
    public void remove() {
        super.remove();
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.remove(this);
    }

    public static Map<ThreadLocal<?>, Object> getTlMap() {
        return THREAD_LOCAL_MAP.get();
    }

    public static void clear(Thread mainThread) {
        if (!Objects.equals(Thread.currentThread(), mainThread)) {
            Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
            //防止迭代器遍历过程中删除元素导致ConcurrentModificationException
            List<Map.Entry<ThreadLocal<?>, Object>> list = new ArrayList<>(map.entrySet());
            list.forEach(entry -> entry.getKey().remove());
            map.clear();
            THREAD_LOCAL_MAP.remove();
        }
    }
}

当然这里还有个坑，你会发现在调用TransferableThreadLocal的clear方法时传了一个线程实例进来，这是为了防止一种场景：当线程池的拒绝策略是ThreadPoolExecutor.CallerRunsPolicy()时，由于执行任务的是当前主线程而非线程池的线程，因此在执行完任务后会把主线程的所有ThreadLocal值清除，后续想要拿ThreadLocal值就会拿不到了，因此才会在创建Runnable时把当前主线程实例的引用传入，防止这种场景下主线程的ThreadLocal值被清除。主线程你就自己逐个去remove，或者重载个clear方法不带if条件判断在主线程执行完后调用就行了。（之前考虑不周请谅解）

4.验证

1.功能验证

我们分别使用InheritableThreadLocal和TransferableThreadLocal验证：

可以看到因为线程复用，InheritableThreadLocal使用的还是原来第一次创建线程时设置的值（没做remove操作），而TransferableThreadLocal每次都能取到最新设置的值。

2.性能测试

性能测试主要分为两块，一是基础功能的测试，二是异步任务的测试。测试比较的对象以阿里提供的TransmittableThreadLocal来做相应的比较，其原理也和上述实现类似，可以参考：GitHub - alibaba/transmittable-thread-local: 📌 TransmittableThreadLocal (TTL), the missing Java™ std lib(simple & 0-dependency) for framework/middleware, provide an enhanced InheritableThreadLocal that transmits values between threads even using thread pooling components.

环境：JDK1.8，CPU：12核，测试使用JMH进行测试：

1.基础功能性能测试

 基础功能主要包括set、get、remove方法的测试，测试的模式包括吞吐量和平均响应时间，考虑到生产服务器分配的是4核CPU，因此采用4个线程进行测试，测试代码如下：

@BenchmarkMode({Mode.Throughput, Mode.AverageTime})
@Warmup(iterations = 5, time = 5)
@Measurement(iterations = 5, time = 5)
@Threads(4)
@Fork(1)
@State(value = Scope.Benchmark)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public class TtlTest {

    private static final ThreadLocal<Object> THREAD_LOCAL_1 = new TransferableThreadLocal<>();
    private static final ThreadLocal<Object> THREAD_LOCAL_2 = new TransmittableThreadLocal<>();

    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(TtlTest.class.getSimpleName())
                .result("result.json")
                .resultFormat(ResultFormatType.JSON).build();
        new Runner(opt).run();
    }

    @Benchmark
    public void TransferableTL() {
        try {
            THREAD_LOCAL_1.set(666);
            THREAD_LOCAL_1.get();
        } finally {
            THREAD_LOCAL_1.remove();
        }
    }

    @Benchmark
    public void TransmittableTL() {
        try {
            THREAD_LOCAL_2.set(666);
            THREAD_LOCAL_2.get();
        } finally {
            THREAD_LOCAL_2.remove();
        }
    }
}

性能测试结果通过JMH可视化界面展示如下，TransferableThreadLocal的TPS是0.07 ops/ns，平均响应时间是64 ns/op；阿里TransmittableThreadLocal的TPS是0.06 ops/ns，平均响应时间是81 ns/op。

2.异步任务性能测试

异步任务其实就是对包装的Runnable类进行的测试，省略掉线程执行的业务代码，最终就是在包装的Runnable任务内对set、get、remove进行测试，测试的模式包括吞吐量和平均响应时间，考虑到生产服务器分配的是4核CPU，因此采用4个线程进行测试，测试代码如下：

@BenchmarkMode({Mode.Throughput, Mode.AverageTime})
@Warmup(iterations = 5, time = 5)
@Measurement(iterations = 5, time = 5)
@Threads(4)
@Fork(1)
@State(value = Scope.Benchmark)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public class TtlTest2 {

    public static final ThreadLocal<Object> t1 = new TransferableThreadLocal<>();

    public static final ThreadLocal<Object> t2 = new TransmittableThreadLocal<>();

    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(TtlTest2.class.getSimpleName())
                .result("result.json")
                .resultFormat(ResultFormatType.JSON).build();
        new Runner(opt).run();
    }

    @Benchmark
    public void TransferableTL() {
        t1.set(666);
        Runnable runnable = ThreadLocalTransferDecorator.decorate(t1::get);
        runnable.run();
    }

    @Benchmark
    public void TransmittableTL() {
        t2.set(666);
        Runnable runnable = TtlRunnable.get(t2::get, true);
        runnable.run();
    }
}

性能测试结果通过JMH可视化界面展示如下，TransferableThreadLocal的TPS是0.016 ops/ns，平均响应时间是206 ns/op；阿里TransmittableThreadLocal的TPS是0.006 ops/ns，平均响应时间是628 ns/op。

5.总结

JDK提供了本地线程共享的ThreadLocal和在创建线程时可继承的InheritableThreadLocal，但是由于大多数情况都推荐使用线程池来进行异步处理，InheritableThreadLocal就达不到我们的预期。但是可以参考它的实现思路，在进行线程间的传递时（如创建线程或创建Runnable实现类），将与父线程绑定的ThreadLocal值传递给子线程，就可实现父子线程ThreadLocal的传递。可使用上述手动实现的TransferableThreadLocal，也可使用阿里提供的TransmittableThreadLocal，两者的性能都是纳秒级别的操作，对程序性能的影响微乎其微。当然TransferableThreadLocal相较TransmittableThreadLocal提供的功能更为精简，代码实现不超过100行，性能更佳。