在Spring异步线程池中自动传递上下文，这样写轻松又方便

问题

在我们的日常开发中，可以通过@Async注解，很方便地启动一个异步线程。

比如现在有一个用户注册成功后，发送欢迎邮件的需求，在用户注册成功以后，便可以启动一个异步线程，在这个线程中调用邮件服务给用户发消息。

这样，即使邮件服务出了问题，也不会影响到当前用户的注册体验。

问题在于，异步线程无法获取原线程的数据信息，如果每次通过手写参数传递又会比较麻烦，所以我们希望通过某种形式，让数据可以自动传递给子线程。

解决方案

1，新建一个类，重写TaskDecorator类的decorate的方法

public class MDCContextDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        //RequestAttributes context = RequestContextHolder.currentRequestAttributes();
        //这里获取是mdc的上下文，也可以获取RequestContextHolder，具体根据你的业务需要操作即可
        Map<String,String> previous = MDC.getCopyOfContextMap();
        return () -> {
            try {
                if (previous != null) {
                    MDC.setContextMap(previous);
                }
                runnable.run();
            } finally {
              //务必记得clear，否则可能会产生内存泄露
                MDC.clear();
            }
        };
    }
}

2，在自定义的线程池中设置我们的自定义装饰器

    @Bean
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        // 设置核心线程数
        executor.setCorePoolSize(20);
        // 设置最大线程数
        executor.setMaxPoolSize(30);
        // 设置队列容量
        executor.setQueueCapacity(1000);
        // 设置线程活跃时间（秒）
        executor.setKeepAliveSeconds(60);
        // 设置默认线程名称
        executor.setThreadNamePrefix("job-");
        // 设置拒绝策略
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        // 等待所有任务结束后再关闭线程池  全局线程池不能关闭
        executor.setWaitForTasksToCompleteOnShutdown(true);
        //设置我们自定义的Decorator
        executor.setTaskDecorator(new MDCContextDecorator());
        return executor;
    }

原理探究

Spring给我们预留一个任务装饰器TaskDecorator，通过这个任务装饰器，可以像AOP一样，对线程做一些功能增强。

在ThreadPoolTaskExecutor的源码中，initializeExecutor方法对线程池进行初始化，会判断是否有装饰器的实现。

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;
	}

总结

利用ThreadPoolTaskExecutor的TaskDecorator，动态的给一个对象添加一些额外的功能，比生成子类会更加灵活。在我们平常的编码过程中，也建议大家尝试使用装饰模式优化我们的代码。