线程池任务装饰器

这是我参与2022首次更文挑战的第7天，活动详情查看：2022首次更文挑战

一、背景

若使用过线程池进行任务处理过就会知道，如果业务线程将具体的执行任务提交到线程池中的线程进行处理，那么如果具体的任务执行需要获取到原业务线程的上下文信息，这时该如何处理？总不能每次手动将原业务线程中的上下文信息作为调用参数，传递给线程池的线程中。在阅读线程池相关信息后发现，线程池中提供了一个装饰器的功能，刚好可以简单有效的处理这一场景。

二、正文

创建线程池装饰器步骤：

创建一个类继承于TaskDecorator接口，并实现它的decorate方法，返回原任务的包装任务。

/**
 * 上下文复制装饰器，将上下文传递给线程池的每一个线程任务
 * 
 * @Author: xiaocainiaoya
 * @Date: 2022/01/28 19:00:19
 **/
@Slf4j
public class ContextCopyingTaskDecorator implements TaskDecorator {

    /**
     * 具体上下文拷贝实现
     * 
     * @Author: xiaocainiaoya
     * @Date: 2022/01/28 19:00:58
     * @param runnable
     * @return:  
     **/
    @Override
    public Runnable decorate(Runnable runnable) {
        //获取父线程上下文信息
        Map<String, Object> context = ThreadContextHandler.getThreadLocal();
        return () -> {
            try {
                //复制到子线程
                ThreadContextHandler.set(context);
                runnable.run();
            } finally {
              // 清理线程上下文
              ThreadContextHandler.clear();
            }
        };
    }
}
复制代码

在线程池创建时将上述类的对象实例化通过taskExecutor.setTaskDecorator添加到创建配置中。

@Configuration
public class ThreadPoolConfig {
    @Bean("serverTaskExecutor")
    public ThreadPoolTaskExecutor taskExecutor() {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        // 处理器的核心数
        int coreNum = Runtime.getRuntime().availableProcessors() * 2;
        // IO密集
        // 设置最大线程池线程数量为核心线程池的两倍
        int maxPoolSize = coreNum * 2;
        taskExecutor.setCorePoolSize(coreNum);
        taskExecutor.setMaxPoolSize(maxPoolSize);
        taskExecutor.setQueueCapacity(maxPoolSize);
        taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        taskExecutor.setThreadNamePrefix("serverTask-");
        // 添加任务装饰器
        taskExecutor.setTaskDecorator(new ContextCopyingTaskDecorator());
        taskExecutor.initialize();
        return taskExecutor;
    }
}
复制代码

在线程池的初始化时，会判断当前线程池是否存在taskDecorator装饰器，若存在则将该装饰器的处理添加到自身的execute方法中（自身也实现于Runnable接口），也就是当线程调用时会先执行该方法，那么就可以在此进行线程上下文信息的复制。

三、问题

在使用了一段时间后发现了一些问题，由于线程创建时配置的拒绝策略为ThreadPoolExecutor.CallerRunsPolicy()，当请求线程过多导致队列满时，会使用调用的业务线程进行业务逻辑的处理，而这里的finally中在线程结束后会清理上下文信息，这就导致了，当并发达到一定程度时，由于原业务线程被当成线程池线程进行业务处理，而导致执行结束后，上下文信息被清理，使得后面的业务执行中获取不到对应的上下文数据信息。

通过日志发现，原业务线程的线程名称都是由http-nio-xxx，所以简单处理就是对名称进行匹配，若匹配成功则不清理上下文信息。

@Slf4j
public class ContextCopyingTaskDecorator implements TaskDecorator {

    static final String HTTP_START_NAME = "http-nio";

    /**
     *   线程池拒绝策略为:CallerRunsPolicy. 当请求线程过多导致队列满时会使用调用线程进行处理,
     *   调用线程不进行线程上下文复制及清除逻辑。
     *
     * @Author: xiaocainiaoya
     * @Date: 2022/01/28 19:48:14
     * @param runnable
     * @return:
     **/
    @Override
    public Runnable decorate(Runnable runnable) {
        //获取父线程上下文信息
        Map<String, Object> context = ThreadContextHandler.getThreadLocal();
        return () -> {
            // 线程名称不是以http-nio开头的线程, 则为线程池线程
            boolean threadPoolThread = !Thread.currentThread().getName().startsWith(HTTP_START_NAME);
            try {
                if(threadPoolThread){
                    log.info("线程任务传递线程上下文信息: {}", JSONObject.toJSONString(context));
                    //复制到子线程
                    ThreadContextHandler.set(context);
                }
                runnable.run();
            } finally {
                if(threadPoolThread){
                    ThreadContextHandler.clear();
                }
            }
        };
    }
}
复制代码