背景说明

在一个分布式系统中，由于链路较多，需要使用一个traceId将整个请求链路串联起来，便于记录和排查问题。
针对每个请求，能够完整的将整个生命周期内的日志记录并且串联起来，无需在每个API中显示（侵入式）编写记录traceId逻辑， 并且针对异步线程，定时任务等，也能生成唯一的traceId跟踪；

解决方案

Slf4j 的MDC (Mapped Diagnostic Context），是基于ThreadLocal实现，能够将特定的信息（如请求ID、用户ID等）与当前线程绑定，使得在多线程环境中生成的日志能够包含这些上下文信息，从而方便追踪和分析问题， 完成全链路traceId集成。

Sping Boot 中的HandlerInterceptor接口，可以用于在请求被DispatcherServlet处理之前或之后，提供一种机制来执行一些自定义的逻辑

HandlerInterceptor接口定义了三个方法：

1. preHandle(HttpServletRequest request, HttpServletResponse response, Object handler): 在请求被处理之前调用。如果该方法返回false，则请求将不会继续执行，相当于一个过滤器。如果返回true，则请求会正常执行。
2. postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView): 在请求被处理之后，DispatcherServlet视图渲染之前调用。可以在这个方法中对ModelAndView对象进行修改。
3. afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex): 请求处理完成，视图渲染之后调用。可以用来进行日志记录、资源清理等工作。

实现逻辑

实现拦截器

新建一个TraceIdInterceptor 需要实现 HandlerInterceptor 接口

@Component
@Slf4j
public class TraceIdInterceptor  implements HandlerInterceptor {

    /**
     * preHandle为请求前拦截
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        // 在日志框架中的MDC中添加请求的唯一标识
        String traceId = String.valueOf(UUID.randomUUID());

        // 绑定key值到线程中
        MDC.put("traceId", traceId);

        // 继续执行接口请求
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, @Nullable Exception ex) {
        MDC.clear();
    }
}

注册拦截器

@Slf4j
@Configuration
public class Interceptor implements WebMvcConfigurer{
    /**
     * http请求拦截器
     * @param registry
     */
    @Override
    public void addInterceptors(InterceptorRegistry registry) {

        registry.addInterceptor(new TraceIdInterceptor());
    }
}

修改response

利用 ResponseBodyAdvice 接口，在响应体写入HTTP响应之前进行自定义处理。

ResponseBodyAdvice接口定义了以下方法：

1. beforeBodyWrite(BodyAdviceContext context, @RequestBody Object body, HttpHeaders headers): 在响应体写入之前调用，可以修改响应体body和HTTP响应头headers。
2. afterBodyWrite(BodyAdviceContext context, @RequestBody Object body, HttpHeaders headers): 在响应体写入之后调用，可以进行一些清理工作。

ResponseBodyAdvice接口的实现类需要添加@ControllerAdvice或@RestControllerAdvice注解，这样Spring容器才能识别并注册这个组件。

修改上述的Interceptor:

@Slf4j
@Configuration
@ControllerAdvice
public class Interceptor implements WebMvcConfigurer, ResponseBodyAdvice {
    /**
     * http请求拦截器
     * @param registry
     */
    @Override
    public void addInterceptors(InterceptorRegistry registry) {

        registry.addInterceptor(new TraceIdInterceptor());
    }

    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body , MethodParameter returnType, MediaType selectedContentType, Class selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        if (body instanceof Response entity) {
            String traceId = MDC.get("traceId");
            entity.setTraceId(traceId);
        }

        return body;
    }
}

修改log配置

修改 resources/logback.xml 配置， 在 <encoder> 中增加

[%X{traceId}]

即可以在日志中记录 traceId

处理异步线程

如果使用了线程池, 需要将本线程的 traceId 传递到 新线程中去，不然新老线程的日志就没办法通过 traceId 关联起来。

采用 AsyncConfigurerSupport + TaskDecorator 实现对异步线程的traceId传递， 参考代码如下:

public class MdcTaskDecorator implements TaskDecorator {
    /**
     * 使异步线程池获得主线程的上下文
     * @param runnable
     * @return
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        Map<String,String> map = MDC.getCopyOfContextMap();
        return () -> {
            try{
                MDC.setContextMap(map);
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }
}

@Configuration
public class AsyncConfig extends AsyncConfigurerSupport {

    @Override
    public Executor getAsyncExecutor() {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        threadPoolTaskExecutor.setCorePoolSize(16);
        threadPoolTaskExecutor.setMaxPoolSize(48);
        threadPoolTaskExecutor.setQueueCapacity(16);
        threadPoolTaskExecutor.setThreadNamePrefix("Async-");
        threadPoolTaskExecutor.setTaskDecorator(new MdcTaskDecorator());
        threadPoolTaskExecutor.initialize();
        return threadPoolTaskExecutor;
    }

}

针对@Scheduled

基于AOP切面，实现对schedule的拦截处理，参考示例如下:

@Aspect
@Component
public class ScheduledTaskAspect {


    // 定义一个切点，匹配所有使用 @Scheduled 注解的方法
    @Pointcut("@annotation(org.springframework.scheduling.annotation.Scheduled)")
    public void scheduledTaskPointcut() {}

    // 前置通知
    @Before("scheduledTaskPointcut()")
    public void beforeScheduledTask(JoinPoint joinPoint) {
        // 前置逻辑，例如记录任务开始日志
        String traceId = String.valueOf(UUID.randomUUID());

        // 绑定key值到线程中
        MDC.put("traceId", traceId);
    }

    @After("scheduledTaskPointcut()")
    public void afterScheduledTask(JoinPoint joinPoint) {
        // 后置逻辑，例如记录任务结束日志
        MDC.clear();
    }

    @AfterThrowing(pointcut = "scheduledTaskPointcut()", throwing = "error")
    public void afterScheduledTaskWithException(JoinPoint joinPoint, Throwable error) {
        // 异常处理逻辑
        MDC.clear();
    }
}

测试验证

至此，针对API , 异步线程，定时任务，我们都可以记录traceId