线程池一切换 traceId 就丢,我在 gRPC 的拦截器和异步调用链上做了三层上下文固定

gRPC 的链路追踪最烦人的不是怎么埋点,而是 traceId 在线程池切换、异步调用时莫名其妙就丢了。我在一个网关项目里踩过这个坑,请求进来时拦截器明明把 traceId 写进了上下文,到了业务线程池里一打印日志,traceId 变 null 了。最后靠三层上下文固定机制才彻底解决:gRPC 拦截器提取元数据 → Context 跨线程传递 → 异步回调再绑定。

下面把完整的实现方案拆开讲,代码可以直接拿去用。

第一层:拦截器里把元数据灌进 Context

gRPC 的元数据通过 Metadata 传过来,拦截器是第一个能接触到请求的地方。服务端拦截器实现 ServerInterceptor,从 Metadata 里把 traceId 和用户信息拽出来,塞进 gRPC 的 Context 里。

public class TraceServerInterceptor implements ServerInterceptor {
    private static final Metadata.Key<String> TRACE_ID_KEY =
            Metadata.Key.of("x-trace-id", Metadata.ASCII_STRING_MARSHALLER);
    private static final Metadata.Key<String> USER_ID_KEY =
            Metadata.Key.of("x-user-id", Metadata.ASCII_STRING_MARSHALLER);
    private static final Context.Key<String> CONTEXT_TRACE_ID =
            Context.key("traceId");
    private static final Context.Key<String> CONTEXT_USER_ID =
            Context.key("userId");

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
            ServerCall<ReqT, RespT> call,
            Metadata headers,
            ServerCallHandler<ReqT, RespT> next) {

        String traceId = headers.get(TRACE_ID_KEY);
        String userId = headers.get(USER_ID_KEY);

        if (traceId == null) {
            traceId = generateTraceId(); // UUID.randomUUID().toString().replace("-", "")
        }

        Context ctx = Context.current()
                .withValue(CONTEXT_TRACE_ID, traceId)
                .withValue(CONTEXT_USER_ID, userId);

        return Contexts.interceptCall(ctx, call, headers, next);
    }
}

这里有个细节:Contexts.interceptCall() 会让后续所有在当前线程上执行的 gRPC 回调都跑在这个 Context 里。但问题也藏在这——一旦你把业务逻辑丢给线程池,gRPC 的 Context 就断链了。gRPC 的 Context 是基于 ThreadLocal 的,线程一切换,它不跟你走。

第二层:线程池切换时显式传递 Context

业务代码几乎不可能全跑在 gRPC 的 IO 线程上。我们用的是自定义线程池,提交任务时必须把 Context 从当前线程"搬运"到工作线程。gRPC 提供了 Context.wrap() 方法来做这件事。

public class ContextAwareExecutor {
    private final ExecutorService delegate;

    public ContextAwareExecutor(ExecutorService delegate) {
        this.delegate = delegate;
    }

    public void execute(Runnable task) {
        Context ctx = Context.current();
        delegate.execute(ctx.wrap(task));
    }

    public <T> Future<T> submit(Callable<T> task) {
        Context ctx = Context.current();
        return delegate.submit(ctx.wrap(task));
    }
}

用的时候,业务代码通过 Context.key() 定义的静态 Key 取值,而不是从 ThreadLocal 或 MDC 里取:

public class TraceContext {
    private static final Context.Key<String> TRACE_ID = Context.key("traceId");
    private static final Context.Key<String> USER_ID = Context.key("userId");

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

    public static String getUserId() {
        return USER_ID.get();
    }
}

实际业务方法里直接调用 TraceContext.getTraceId() 就能拿到当前请求的 traceId,不管跑在哪个线程上。我自己的习惯是在每个 @Service 方法的入口打一行日志,确认 traceId 还在:

log.info("[traceId={}] 开始处理订单, userId={}", TraceContext.getTraceId(), TraceContext.getUserId());

到这步,同步调用链路的问题基本解决了。但异步回调才是真正的重灾区。

第三层:异步回调链上的 Context 再绑定

gRPC 服务端返回响应时经常会异步处理,比如调用下游服务后通过 StreamObserver.onNext() 返回结果。问题在于,当你用 CompletableFuture 做异步编排时,.thenApply().thenAccept() 里的回调线程是 ForkJoinPool 的公共线程,Context 又丢了。

演示一下错误写法:

// 错误:thenApplyAsync 里的线程拿不到 Context
CompletableFuture.supplyAsync(() -> queryOrder())
    .thenApplyAsync(order -> {
        // 这里 TraceContext.getTraceId() 返回 null
        log.info("处理订单");
        return enrichOrder(order);
    });

正确做法是每次异步回调都重新绑定当前 Context:

public class ContextAsyncUtil {
    public static <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
        Context ctx = Context.current();
        return CompletableFuture.supplyAsync(() -> {
            Context previous = ctx.attach();
            try {
                return supplier.get();
            } finally {
                Context.current().detach(previous);
            }
        });
    }

    public static <T, U> CompletableFuture<U> thenApplyAsync(
            CompletableFuture<T> future,
            Function<T, U> fn) {
        Context ctx = Context.current();
        return future.thenApplyAsync(result -> {
            Context previous = ctx.attach();
            try {
                return fn.apply(result);
            } finally {
                Context.current().detach(previous);
            }
        });
    }
}

ctx.attach() 会把 Context 挂到当前线程上并返回原来的 Context,detach() 恢复回去。这个模式是 gRPC 官方文档推荐的,但在实际项目里很容易被忽略。

业务代码改成这样:

CompletableFuture<Order> orderFuture = ContextAsyncUtil.supplyAsync(() -> queryOrder());

ContextAsyncUtil.thenApplyAsync(orderFuture, order -> {
    log.info("[traceId={}] 订单查询完成", TraceContext.getTraceId());
    return enrichOrder(order);
});

如果你的项目里用了 Spring 的 @Async,也需要做类似的处理。自定义一个 AsyncConfigurer,把 Context 传递进去:

@Configuration
public class AsyncConfig implements AsyncConfigurer {
    @Override
    public Executor getAsyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(10);
        executor.setMaxPoolSize(50);
        executor.setTaskDecorator(task -> {
            Context ctx = Context.current();
            return ctx.wrap(task);
        });
        executor.initialize();
        return executor;
    }
}

TaskDecorator 是 Spring 4.3 之后提供的钩子,可以在每个任务提交前做一层包装,正好用来绑定 Context。

日志输出:把 Context 值打进 MDC

光在代码里能拿到 traceId 还不够,日志里也得自动带上,不然查问题还得人肉 grep。用 SLF4J 的 MDC,在拦截器入口放进去,请求结束时清掉:

// 在 ServerInterceptor 的 interceptCall 方法里
String traceId = headers.get(TRACE_ID_KEY);
if (traceId == null) {
    traceId = generateTraceId();
}
MDC.put("traceId", traceId);
MDC.put("userId", userId);

// 在 ServerCall.Listener 的 onComplete 或 onCancel 回调里清理
@Override
public void onComplete() {
    MDC.clear();
    super.onComplete();
}

但注意,MDC 也是基于 ThreadLocal 的,线程池切换时同样会丢。所以前面第二层、第三层的 Context 绑定搞好了,MDC 的值才能跟着走。在 Context.wrap() 的 Runnable 里可以加一行:

public void execute(Runnable task) {
    Context ctx = Context.current();
    Map<String, String> contextMap = MDC.getCopyOfContextMap();
    delegate.execute(() -> {
        if (contextMap != null) {
            MDC.setContextMap(contextMap);
        }
        ctx.wrap(task).run();
        MDC.clear();
    });
}

这样日志格式里配置 %X{traceId} 就能在每个线程的日志行里打出正确的 traceId。

客户端拦截器:出站请求也要带上 traceId

服务端能收到 traceId 的前提是客户端发请求时把它塞进了 Metadata。客户端拦截器实现 ClientInterceptor

public class TraceClientInterceptor implements ClientInterceptor {
    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
            MethodDescriptor<ReqT, RespT> method,
            CallOptions callOptions,
            Channel next) {

        return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
                next.newCall(method, callOptions)) {

            @Override
            public void start(Listener<RespT> responseListener, Metadata headers) {
                String traceId = TraceContext.getTraceId();
                if (traceId != null) {
                    headers.put(
                        Metadata.Key.of("x-trace-id", Metadata.ASCII_STRING_MARSHALLER),
                        traceId
                    );
                }
                super.start(responseListener, headers);
            }
        };
    }
}

这里直接从 TraceContext.getTraceId() 取值,前提是客户端调用链也经过了 Context 绑定。如果客户端是独立进程(比如一个定时任务触发 gRPC 调用),那就自己生成一个 traceId 放进去。

常见问题

gRPC 的 Context 和 ThreadLocal 到底有什么区别?

gRPC 的 Context 是一个不可变的、支持取消和超时的键值存储,它的设计目标就是跨线程传播。ThreadLocal 只在当前线程可见,gRPC 的 Context 可以通过 attach()/detach() 显式绑定到任意线程,也可以通过 wrap() 把 Runnable/Callable 包一层自动传递。但 Context 本身不会主动跟着线程走,必须手动处理——这也是为什么线程池切换时容易丢。

为什么不直接用 TransmittableThreadLocal 或者阿里的 TTL?

可以用,但没必要引入额外依赖。gRPC 的 Context 机制本身就解决了这个问题,而且和 gRPC 的生命周期管理(取消信号、deadline)深度集成。如果你的项目里有非 gRPC 的异步链路(比如 MQ 消费、定时任务),TTL 确实更通用,但对于纯 gRPC 调用链,Context.wrap() 足够。

MDC 在 CompletableFuture 回调里丢了怎么办?

和 Context 丢的原因一样——线程切换了。解决方案就是在 ContextAsyncUtil 的包装方法里同时传递 MDC,或者在 ctx.attach() 之后手动 MDC.setContextMap()。也可以考虑用 logback-mdc-ttl 这种现成的适配库,但核心思路不变:每次异步回调都要显式传递。