Spring AI：Reactor 异步执行中的线程上下文传递实践

背景问题

在使用 Spring AI MCP Server 框架时，工具方法可能被框架在异步线程中执行。当线程切换时，传统的 ThreadLocal 无法传递上下文信息（如 token、traceId），导致：

认证失败：AppThreadLocal.getToken() 返回 null
链路追踪中断：AppThreadLocal.getTraceId() 返回 null
日志无法关联：无法追踪完整的请求链路

核心解决方案

方案概述

通过替换 Reactor 的全局调度器工厂 ，使用 TTL（Transmittable Thread Local）包装的线程池，确保线程上下文在异步执行时自动传递。

核心步骤

步骤 1：引入依赖

复制代码

<!-- Reactor 核心库 -->
<dependency>
    <groupId>io.projectreactor</groupId>
    <artifactId>reactor-core</artifactId>
    <version>3.7.3</version>
</dependency>

<!-- TTL 线程本地变量传递 -->
<dependency>
    <groupId>com.alibaba</groupId>
    <artifactId>transmittable-thread-local</artifactId>
    <version>2.11.5</version>
</dependency>

步骤 2：使用 TransmittableThreadLocal 存储上下文

复制代码

public class AppThreadLocal {
    // 使用 TransmittableThreadLocal 而非普通 ThreadLocal
    private static final ThreadLocal<HeaderModel> REQUEST_CONTEXT =
            new TransmittableThreadLocal<>();

    public static void setToken(String token) {
        HeaderModel header = get();
        if (header == null) {
            header = new HeaderModel();
        }
        header.setToken(token);
        set(header);
    }

    public static String getToken() {
        HeaderModel header = get();
        return header != null ? header.getToken() : null;
    }
}

关键点 ：TransmittableThreadLocal 支持在线程池间传递，但需要配合 TTL 包装的线程池使用。

步骤 3：实现自定义 Scheduler

复制代码

public class MyCustomBoundedElasticScheduler implements Scheduler {
    private final ExecutorService executorService;
    private final String name;

    public MyCustomBoundedElasticScheduler(ExecutorService executorService, String name) {
        this.executorService = executorService;
        this.name = name;
    }

    @Override
    public Disposable schedule(Runnable task) {
        // 任务提交到 TTL 包装的线程池
        executorService.execute(task);
        return () -> {};
    }

    @Override
    public Worker createWorker() {
        return new ExecutorServiceWorker(executorService);
    }

    // ... 其他方法实现
}

步骤 4：替换全局调度器工厂

复制代码

@Configuration
public class ReactorConfig {

    @Bean
    public Object replaceReactorSchedulerFactory() {
        // 替换 Reactor 的全局调度器工厂
        Schedulers.setFactory(new Schedulers.Factory() {
            @Override
            public Scheduler newBoundedElastic(
                    int threadCap, 
                    int queuedTaskCap, 
                    ThreadFactory threadFactory, 
                    int ttlSeconds) {
                
                // 关键：使用 TTL 包装线程池
                ExecutorService ttlExecutor = TtlExecutors.getTtlExecutorService(
                    Executors.newFixedThreadPool(50)
                );
                
                return new MyCustomBoundedElasticScheduler(
                    ttlExecutor, 
                    "newBoundedElastic"
                );
            }
        });
        
        return new Object();
    }
}

工作原理：

当框架调用 Schedulers.newBoundedElastic() 时
触发自定义工厂方法
返回使用 TTL 包装线程池的自定义 Scheduler
任务执行时，TTL 自动传递线程上下文

步骤 5：在请求入口设置上下文

复制代码

@Slf4j
public class WsFilter extends GenericFilter {
    
    public void doFilter(ServletRequest request, ServletResponse response, 
                        FilterChain chain) {
        if (request.getRequestURI().equals("/mcp/message")) {
            // 从请求中提取上下文信息
            String token = extractToken(request);
            String traceId = extractTraceId(request);
            
            // 设置到 TransmittableThreadLocal
            AppThreadLocal.setToken(token);
            AppThreadLocal.setTraceId(traceId);
        }
        
        chain.doFilter(request, response);
    }
}

步骤 6：在工具方法中使用上下文

复制代码

@Service
public class McsService extends BaseService {
    
    @Tool(name = "queryLoginCompanyBasicSkuPageList")
    public String queryLoginCompanyBasicSkuPageList() {
        // 即使在线程池线程中执行，也能获取到上下文
        String token = AppThreadLocal.getToken();      
        String traceId = AppThreadLocal.getTraceId();  
        
        return execute("ech-mcs/v1/shop/sku/...", HttpMethod.POST, param, query);
    }
}

执行流程

复制代码

请求到达 → WsFilter 设置上下文 → 框架异步执行工具
    ↓
框架调用 Schedulers.newBoundedElastic()
    ↓
触发自定义工厂 → 返回 TTL 包装的 Scheduler
    ↓
任务提交到 TTL 线程池 → 自动捕获当前线程上下文
    ↓
在新线程中执行 → 自动恢复上下文
    ↓
工具方法获取上下文 ✅ → 正常执行