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();
    }
}

工作原理：

1. 当框架调用 Schedulers.newBoundedElastic() 时

2. 触发自定义工厂方法

3. 返回使用 TTL 包装线程池的自定义 Scheduler

4. 任务执行时，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 线程池 → 自动捕获当前线程上下文
    ↓
在新线程中执行 → 自动恢复上下文
    ↓
工具方法获取上下文 ✅ → 正常执行