优雅的解决异步线程TraceId传递问题

1. 问题

在业务开发过程中，我们通常需要对整个请求链路进行追踪，通过日志追踪的方式定位链路上的问题点，那么主要采取的方式就是在请求中传递TraceId，同一份请求链路TraceId唯一，通过在日志或监控获取请求TraceId来观测整个链路。

如果我们业务开发中使用了异步方式处理业务逻辑，比如使用异步方式调用下游接口，如果不做特殊处理的话TraceId是无法传递给异步线程的，也就是说异步线程无法获取到主线程的上下文信息。本文说明的就是这个问题的优雅解决方法。

2. 异步例子

定义上下文信息，用来存储TraceId信息

public class UserContext {

    private static final ThreadLocal<String> threadLocal = new ThreadLocal<>();

    public static void setTraceId(String traceId) {
        threadLocal.set(traceId);
    }

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

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

定义请求连接器，用来解析header中的TraceId（假设上游通过header传递TraceId）

@Component
public class TraceIdInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String traceId = request.getHeader("traceId");
        UserContext.setTraceId(traceId);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        // 释放资源，防止内存泄漏
        UserContext.remove();
    }
}

定义服务接口，用来接受请求

@RestController
public class TraceIdController {

    @Resource
    private TraceIdService service;

    @GetMapping("/api")
    public String api() {
        String traceId = UserContext.getTraceId();
        System.out.println(traceId);

        service.asyncCall();
        return "hello world";
    }
}

定义service异步处理业务（包括调用下游接口）

@Service
public class TraceIdService {

    /**
     * 异步调用
     */
    @Async("taskExecutor")
    public void asyncCall() {
        String traceId = UserContext.getTraceId();
        System.out.println("traceId: " + traceId);
    }
}

定义taskExecutor线程池，异步任务也走这个线程池

@Configuration
public class ThreadPoolConfig {

    @Bean("taskExecutor")
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(1000);
        executor.setThreadNamePrefix("ThreadPool-");
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.setAwaitTerminationSeconds(60);
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.initialize();
        return executor;
    }
}

接下来我们请求一下看一下输出结果

我们看到异步执行的Service层并没有获取到TraceId的值，这是因为ThreadLocal+异步线程导致的，那怎么解决这个问题呢？

3. 解决问题

我们可以采用ThreadPoolTaskExecutor线程池中的TaskDecorator来解决，任务装饰器可以提供在任务执行前后增加一些处理逻辑，方便我们扩展任务。这个装饰器的目的就是在线程执行之前把主线程的数据传递给子线程。

自定义一个TaskDecorator

public class TraceIdTaskDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        String traceId = UserContext.getTraceId();
        // 日志相关的上下文信息
        Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
        return () -> {
            UserContext.setTraceId(traceId);
            MDC.setContextMap(copyOfContextMap);
            runnable.run();
        };
    }
}

线程池设置TaskDecorator

@Configuration
public class ThreadPoolConfig {

    @Bean("taskExecutor")
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(1000);
        executor.setThreadNamePrefix("ThreadPool-");
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.setAwaitTerminationSeconds(60);
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.setTaskDecorator(new TraceIdTaskDecorator()); // 设置TaskDecorator
        executor.initialize();
        return executor;
    }
}

我们再来请求一下，得到如下结果

异步线程已经获取到了主线程传递过来的traceId