Spring线程池—TaskDecorator线程的装饰（跨线程传递ThreadLocal的方案）

ThreadPoolTaskExecutor在执行线程时，存在一个TaskDecorator配置，可以装饰线程类。

1. 源码分析

源码：org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor#initializeExecutor

@Override protected ExecutorService initializeExecutor(ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {

    BlockingQueue < Runnable > queue = createQueue(this.queueCapacity);
    
    ThreadPoolExecutor executor;
    //配置了线程装饰器
    if (this.taskDecorator != null) {
        executor = new ThreadPoolExecutor(this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS, queue, threadFactory, rejectedExecutionHandler) {@Override public void execute(Runnable command) {
                //装饰线程类
                Runnable decorated = taskDecorator.decorate(command);
                if (decorated != command) {
                    decoratedTaskMap.put(decorated, command);
                }
                //执行原有逻辑
                super.execute(decorated);
            }
        };
    } else {
        executor = new ThreadPoolExecutor(this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS, queue, threadFactory, rejectedExecutionHandler);

    }

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

    this.threadPoolExecutor = executor;
    return executor;
}

2. 实际使用

声明spring线程池时，需要指定TaskDecorator，在装饰代码中获取到ThreadLocal的值。

@Configuration
@EnableAsync
public class ThreadPoolConfig {

    @Bean("asyncServiceExecutor")
    public ThreadPoolTaskExecutor asyncRabbitTimeoutServiceExecutor() {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        //核心线程数
        threadPoolTaskExecutor.setCorePoolSize(5);
        //核心线程若处于闲置状态的话，超过一定的时间(KeepAliveTime)，就会销毁掉。
        threadPoolTaskExecutor.setAllowCoreThreadTimeOut(true);
        //最大线程数
        threadPoolTaskExecutor.setMaxPoolSize(10);
        //配置队列大小
        threadPoolTaskExecutor.setQueueCapacity(300);
        //加入装饰器
        threadPoolTaskExecutor.setTaskDecorator(new ContextCopyingDecorator());
        //配置线程池前缀
        threadPoolTaskExecutor.setThreadNamePrefix("test-log-");
        //拒绝策略:只要线程池未关闭，该策略直接在调用者线程中串行运行被丢弃的任务，显然这样不会真的丢弃任务，但是可能会造成调用者性能急剧下降
        threadPoolTaskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        threadPoolTaskExecutor.initialize();
        return threadPoolTaskExecutor;
    }

    @Slf4j
    static class ContextCopyingDecorator implements TaskDecorator {
        @Nonnull
        @Override
        public Runnable decorate(@Nonnull Runnable runnable) {
            //主流程
            String res = TestMappingController.local.get();
            log.info("装饰前：" + res);
            //子线程逻辑
            return () -> {
                try {
                    //将变量重新放入到run线程中。
                    TestMappingController.local.set(res);
                    log.info("打印日志-开始");
                    runnable.run();
                } finally {
                    log.info("打印日志-结束");
                }
            };
        }
    }
}

业务逻辑：

@PostMapping("test2") 
public void test() {
    //主线程执行
    local.set("主线程设置的ThreadLocal");
    //子线程执行
    executor.execute(() - >{
        String res = local.get();
        log.info("线程:" + res);
    });
}

执行结果：可以看到，在ContextCopyingDecorator中是主线程在调用，所以可以获取到主线程的ThreadLocal信息。

2020-11-17 18:54:25,027 INFO [1747069] [http-nio-8087-exec-5] [] (ThreadPoolConfig.java:53): 装饰前：主线程设置的ThreadLocal
2020-11-17 18:54:25,028 INFO [1747070] [test-log-4] [] (ThreadPoolConfig.java:59): 打印日志-开始
2020-11-17 18:54:25,028 INFO [1747070] [test-log-4] [] (TestMappingController.java:46): 线程:主线程设置的ThreadLocal
2020-11-17 18:54:25,028 INFO [1747070] [test-log-4] [] (ThreadPoolConfig.java:62): 打印日志-结束