MDC实现traceId全局追踪

/**
 * 能线程池中也可以传递 traceId
 */
public class MdcTaskExecutorCustomizer implements TaskExecutorCustomizer {

    @Override
    public void customize(ThreadPoolTaskExecutor taskExecutor) {
        taskExecutor.setTaskDecorator(runnable -> {
            Map context = MDC.getCopyOfContextMap();
            return () -> {
                MDC.setContextMap(context);
                try {
                    runnable.run();
                } finally {
                    MDC.clear();
                }
            };
        });
    }
}

/**
 * MDC handler
 */
public class MdcMvcHandlerInterceptorInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        //如果有上层调用就用上层的ID
        String traceId = request.getHeader(MDCUtil.TRACE_ID);
        if (StrUtil.isEmpty(traceId)) {
            traceId = MDCUtil.getNewTraceId();
        }
        MDCUtil.startTrace(traceId);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        MDCUtil.endTrace();
    }
}

/**
 * ####################Servlet######################
 */
@Configuration
@ConditionalOnClass(DispatcherServlet.class)
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
public class MdcInterceptorConfig implements WebMvcConfigurer {

    @Bean
    public HandlerInterceptor mdcMvcHandlerInterceptorInterceptor() {
        return new MdcMvcHandlerInterceptorInterceptor();
    }

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(mdcMvcHandlerInterceptorInterceptor());
    }
}

public class MdcBeanPostProcessor implements InstantiationAwareBeanPostProcessor, ApplicationContextAware {

    private ApplicationContext applicationContext;

    /**
     * 应用到 所有 ThreadPoolTaskExecutor
     * @param bean
     * @param beanName
     * @return
     * @throws BeansException
     */
    @Override
    public boolean postProcessAfterInstantiation(Object bean, String beanName) throws BeansException {
        if (bean instanceof ThreadPoolTaskExecutor) {
            MdcTaskExecutorCustomizer mdcTaskExecutorCustomizer = new MdcTaskExecutorCustomizer();
            ThreadPoolTaskExecutor executor = (ThreadPoolTaskExecutor) bean;
            mdcTaskExecutorCustomizer.customize(executor);
            return true;
        }
        return InstantiationAwareBeanPostProcessor.super.postProcessAfterInstantiation(bean, beanName);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }
}

/**
 * 自动装配
 */
public class MdcAutoConfiguration {


    @Bean
    public MdcBeanPostProcessor mdcBeanPostProcessor() {
        return new MdcBeanPostProcessor();
    }

    /**
     * ####################reactive######################
     */
    @Bean
    @ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.REACTIVE)
    public MdcWebFilter mdcWebFilter() {
        return new MdcWebFilter();
    }

}

/**
 * MDC gateway 过滤器
 */
public class MdcWebFilter implements WebFilter {

    @Override
    public Mono filter(ServerWebExchange exchange, WebFilterChain chain) {
        String traceId = MDCUtil.getNewTraceId();
        // 1.将traceId传递给微服务
        ServerHttpRequest request = exchange.getRequest().mutate().header("traceId", traceId).build();
        // 2.将traceId设置到slf4j中，日志打印模板配置打印traceId
        MDCUtil.startTrace(traceId);
        return chain.filter(exchange).subscriberContext(Context.of(MDCUtil.TRACE_ID, traceId)).doOnTerminate(MDCUtil::endTrace);
    }
}

/**
 * MDC自动配置类入口
 */
public class MdcTracerImportSelector implements ImportSelector {

    @Override
    public String[] selectImports(AnnotationMetadata importingClassMetadata) {
        return new String[]{"com.senge.ali.common.config.mdc.MdcAutoConfiguration"};
    }

}

/**
 * MDC traceId 自定义注解
 */
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Import(MdcTracerImportSelector.class)
public @interface EnabledTraceId {

}

/**
 * MDC traceId 工具类
 */
public abstract class MdcUtil {

    public static final String TRACE_ID = "traceId";

    public static String getNewTraceId() {
        return IdUtil.getSnowflakeNextIdStr();
    }

    public static void startTrace() {
        MDC.put(TRACE_ID, IdUtil.getSnowflakeNextIdStr());
    }

    public static void startTrace(String traceId) {
        MDC.put(TRACE_ID, traceId);
    }

    public static String getCurrentId(){
        return MDC.get(TRACE_ID);
    }

    public static void endTrace() {
        MDC.clear();
    }

}

启动类加注解即可：@EnabledTraceId

@EnabledTraceId
@EnableDiscoveryClient
@SpringBootApplication
public class GatewayApplication {

    public static void main(String[] args) {
        SpringApplication.run(GatewayApplication.class, args);
    }

}