通过TaskDecorator进行线程池上下文传递

一、背景

  需要在线程池的子线程中获取主线程上下文的参数

二、方案

使用spring实现的ThreadPoolTaskExecutor，其中有TaskDecorator装饰器可以实现对线程调用时的加工。

1.线程池定义

为了比较分两种，其一为java原生ThreadPoolExecutor，另一种为spring实现的ThreadPoolTaskExecutor

public class ThreadUtil { 

//java原生ThreadPoolExecutor
 public enum Pool {

       
        COMMON(128);
       

        ExecutorService executor;

        Pool(int threadNum) {
            executor = new ThreadPoolExecutor(threadNum,threadNum,0L,TimeUnit.MILLISECONDS,new LinkedBlockingQueue(threadNum),new ThreadPoolExecutor.CallerRunsPolicy());

        }


        public ExecutorService getExecutor() {
            return executor;
        }
    }


//spring实现的ThreadPoolTaskExecutor
    public enum TaskPool {
   
        COMMON(128) ;

        TaskExecutor taskExecutor;

        private final static String threadNamePrefix = "TaskPool-";

        TaskPool(int threadNum) {
            ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
            executor.setCorePoolSize(threadNum);
            executor.setMaxPoolSize(threadNum);
            executor.setQueueCapacity(threadNum);
            executor.setThreadNamePrefix(threadNamePrefix);
            executor.setTaskDecorator(new ProductTaskDecorator());
            executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
            executor.setWaitForTasksToCompleteOnShutdown(true);
            executor.initialize();
            taskExecutor = executor;
        }


        public TaskExecutor getExecutor() {
            return taskExecutor;
        }
    }
}

2.TaskDecorator装饰器实现

public class ProductTaskDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            Map<String, String> contextMap = new HashMap<>();
            //设置一个随机数用来确保是传递的
            contextMap.put("test", UUID.randomUUID().toString());
            MDC.setContextMap(contextMap);
            return () -> {
                try {
                    MDC.setContextMap(contextMap);
                    runnable.run(); //执行线程方法
                } finally {
                    MDC.clear(); //清除字段对线程池的污染
                }
            };
        } catch (Exception e) {
            return runnable;
        }
    }
}

3.测试上下文传递,伪代码

   Map<String, String> contextMap = new HashMap<>();
        contextMap.put("test","1");
        MDC.setContextMap(contextMap);
CompletableFuture.supplyAsync(() ->
                         xxService.getA(xxx)
                , ThreadUtil.Pool.COMMON.getExecutor());
 
CompletableFuture.supplyAsync(() ->
                         xxService.getB(xxx)
                , ThreadUtil.TaskPool.COMMON.getExecutor());
         
              
        public  Map<String, String> getA(xxx) {log.error("ycx 子线程方法a内 :{}", MDC.get("test"));}
  
        public  Map<String, String> getB(xxx) {log.error("ycx 子线程方法b内 :{}", MDC.get("test"));}

4.结果

[pool-37-thread-192][] ycx test 子线程方法a内 :null
[TaskPool-11][]  ycx test 子线程方法b内 :c34f37e1-fcf4-40ff-aef1-f9b9ef088897

可以看到使用java原生的ThreadPoolTaskExecutor没有获取到值，而通过装饰器加工的能获取到