Spring异步线程池—传递线程上下文(TaskDecorator实现)

问题

在spring中使用@async异步调用的情况下，被调用的异步子线程获取不到父线程的request信息，以便处理相关逻辑，即子线程无法获取父线程的上下文数据

思路

在自定义的异步线程池ThreadPoolTaskExecutor中，初始化线程池时有taskDecorator这样一个任务装饰器，类似aop，可对线程执行方法的始末进行增强。其初始化源码如下

 protected ExecutorService initializeExecutor(ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
        BlockingQueue queue = this.createQueue(this.queueCapacity);
        ThreadPoolExecutor executor;
        if (this.taskDecorator != null) {
            executor = new ThreadPoolExecutor(this.corePoolSize, this.maxPoolSize, (long)this.keepAliveSeconds, TimeUnit.SECONDS, queue, threadFactory, rejectedExecutionHandler) {
                public void execute(Runnable command) {
                    Runnable decorated = ThreadPoolTaskExecutor.this.taskDecorator.decorate(command);
                    if (decorated != command) {
                        ThreadPoolTaskExecutor.this.decoratedTaskMap.put(decorated, command);
                    }

                    super.execute(decorated);
                }
            };
        } else {
            executor = new ThreadPoolExecutor(this.corePoolSize, this.maxPoolSize, (long)this.keepAliveSeconds, TimeUnit.SECONDS, queue, threadFactory, rejectedExecutionHandler);
        }

        if (this.allowCoreThreadTimeOut) {
            executor.allowCoreThreadTimeOut(true);
        }

        this.threadPoolExecutor = executor;
        return executor;
    }

基本使用，自定义装饰器实现TaskDecorator ，重写decorate方法，自定义线程池，并设置自定义装饰器

自定义异步线程池
@Bean("taskExecutor") // bean 的名称，默认为首字母小写的方法名
  public Executor taskExecutor() {
    ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
    //其他参数省略
    //设置装饰器
       threadPoolTaskExecutor.setTaskDecorator(new ContextCopyingDecorator());
    return executor;
  }

自定义装饰器，ContextCopyingDecorator，通过try，finally，在子线程执行完后将该线程设置的上下文变量清除

public class ContextCopyingDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            //获取父线程的context
            RequestAttributes context = RequestContextHolder.currentRequestAttributes();
            return () -> {
                try {
                    //将父线程的context设置进子线程里
                    RequestContextHolder.setRequestAttributes(context);
                    //子线程方法执行
                    runnable.run();
                } finally {
                    //清除子线程context
                    RequestContextHolder.resetRequestAttributes();
                }
            };
        } catch (IllegalStateException e) {
            return runnable;
        }
    }
}

存在问题

从父线程取出的RequestContextHolder对象，此为持有线程上下文的request容器，将其设置到子线程中，按道理只要对象还存在强引用，就不会被销毁，但由于RequestContextHolder的特殊性，在父线程销毁的时候，会触发里面的resetRequestAttributes方法（即清除threadLocal里面的信息，即reques中的信息会被清除），此时即使RequestContextHolder这个对象还是存在，子线程也无法继续使用它获取request中的数据了。这也是网上很多文章讲TaskDecorator时没提到的点，真正用起来会发现有时可以有时不行，这个就取决于父子线程哪个先结束了。

完善思路

既然是RequestContextHolder的特殊性，那我们就让绕过他的销毁清除，思路不变，还是继续使用threadLocal来传递我们需要使用到的变量，在父线程装饰前将所需变量取出来，然后在子线程中设置到threadLocal，业务使用的时候从threadLocal中取即可。

改造，自定义threadLocal类（此例子以ua为例子），修改自定义装饰器逻辑

public class ContextCopyingDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            //获取父线程的request的user-agent(示例)
           HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            String ua = request.getHeader("user-agent");
            return () -> {
                try {
                    //将父线程的ua设置进子线程里
                    ThreadLocalData.setUa(ua);
                    //子线程方法执行
                    runnable.run();
                } finally {
                    //清除线程threadLocal的值
                    ThreadLocalData.remove();
                }
            };
        } catch (IllegalStateException e) {
            return runnable;
        }
    }
}

ThreadLocalData

public class ThreadLocalData {
    public static final ThreadLocal threadLocal = new ThreadLocal<>();

    public static String getUa(){
        return threadLocal.get();
    }

    public static void setUa(String ua){
        threadLocal.set(ua);
    }

    public static void remove(){
        threadLocal.remove();
    }
}

至此经测试，一切符合预期

涉及知识点

ThreadLocal，InheritableThreadLocal，TaskDecorator，RequestContextHolder，TransmittableThreadLocal（通过继承InheritableThreadLocal实现，阿里的，推荐）

测试 ThreadLocal，InheritableThreadLocal，TransmittableThreadLocal的区别和使用

1.父线程使用ThreadLocal，子线程创建时不会拥有父类的threadLocal信息
2.父线程使用InheritableThreadLocal，子线程创建时，默认init方法会拿到父类的InheritableThreadLocal信息，这种在线程池/线程复用的情况下，由于init方法只会在初始化时获取父线程的数据，复用的时候也没法再从父线程那里新的InheritableThreadLocal的数据，此种情况下继续使用，很容易出bug（InheritableThreadLocal适用于非线程池和复用线程，单独创建销毁子线程执行的情况）
3.父线程使用TransmittableThreadLocal，子线程创建时拥有父类的TransmittableThreadLocal信息，在线程池/线程复用的情况下不会出现读取到脏数据的情况

总结

• 在异步线程池的情况下，通过ThreadLocal+TaskDecorator一般即可解决遇到的透传问题（方式1）
• 使用阿里的TransmittableThreadLocal，其原理也是对Runnable，Callable，进行装饰（方式2）

参考

Spring线程池—TaskDecorator线程的装饰（跨线程传递ThreadLocal的方案）
(28条消息) TaskDecorator——异步多线程中传递上下文等变量_WannaRunning的博客-CSDN博客