Zipkin链路追踪异步调用问题

思考

链路追踪为我们排查问题、观测服务之间的调用提供了很大的帮助。Zipkin作为链路追踪的基础框架。在工作中我们经常需要使用异步方式去调用一些方法，很常见的是使用@Async注解或者CompletableFuture .supplyAsync()方法去实现异步的调用。而在过程中发现，使用CompletableFuture .supplyAsync()却会导致链路断掉，导致TraceId丢失，无法追踪当时的调用情况。而使用@Async却不会。
使用@Async的链路：

使用CompletableFuture .supplyAsync()的链路：

那么：

• 为什么CompletableFuture .supplyAsync()和@Async两种方式的结果不同呢？
• 我们又如何解决在使用CompletableFuture .supplyAsync()时链路信息能够正确传播？

为了搞清楚上面的问题，我们得先了解链路追踪是如何实现的，Zipkin是我们公司使用的链路追踪框架，首先我们得了解Zipkin的工作原理。

Zipkin的工作原理

Zipkin的工作原理主要围绕“追踪”和“跨度”这两个概念：

• 追踪（Trace）：一个追踪代表了一个请求在系统中的整个生命周期。它从用户发起请求的那一刻开始，直到请求处理完成为止。追踪可以跨越多个服务和进程，帮助我们了解请求的流转路径。
• 跨度（Span）：一个追踪由多个跨度组成，每个跨度代表请求流转过程中某个服务的处理。每个跨度都有一个开始时间和结束时间，还可以包含服务名、操作名、请求的元数据（如 HTTP 方法、状态码等）。每个跨度通常对应一个服务内的操作，多个跨度通过唯一的TraceId串联在一起。

我们首先需要了解Zipkin中几个重要的组件：

• Tracer：
  • Tracer 是分布式追踪系统的核心对象，负责创建、管理和操作追踪信息（Span）。它是追踪框架交互的入口，通常是你通过它来开始和结束链路追踪。
• Span：
  • Span 是一个表示操作的基本单位，通常表示服务中某一具体的操作或请求的生命周期。在一个分布式系统中，Span 通过 TraceId 和 SpanId 关联多个微服务和操作。
• TraceContext：
  • TraceContext 是链路追踪的上下文对象，包含了追踪操作的全局信息，如 TraceId 和 SpanId。TraceContext 是分布式追踪的核心，它能够关联多个 Span，并确保整个调用链中追踪信息的一致性。TraceContext是基于ThreadLocal实现的。
• Scope：
  • Scope 用来管理和维护当前线程的活动上下文，尤其是在多线程或异步操作中，它确保当前线程能够访问到正确的 TraceContext（即 Span 的上下文）。

对象	描述	创建/管理者	关系
Tracer	分布式追踪的核心，负责创建和管理 Span 和 TraceContext。	由应用或框架初始化	创建 Span，并将 Span 与 TraceContext 关联。
Span	代表一次操作的追踪信息，包含操作的起始、结束、元数据等。	Tracer创建	包含 TraceContext，用于标识某个操作的生命周期。
TraceContext	追踪的上下文，包含 TraceId 和 SpanId 等标识信息。	Tracer创建	每个 Span 都持有一个 TraceContext，用于跨服务传递追踪信息。
Scope	表示当前线程的追踪上下文，确保线程或任务中能够关联正确的 TraceContext。	Tracer创建	维护当前线程或任务的 TraceContext，并确保上下文正确传播。

工作流程：

1. 创建链路追踪：
   a. 当请求进入系统时，Tracer 会创建一个新的 Span，并生成相应的 TraceContext（例如，TraceId 和 SpanId）。
   b. Span 和 TraceContext 会通过请求的属性传递到后续的服务或线程中。
2. 传递链路追踪上下文：
   a. 在分布式系统中，TraceContext 会通过 HTTP 请求的 header（例如 X-B3-TraceId 和 X-B3-SpanId）传递给下游服务。
   b. 在异步操作或多线程中，Scope 会确保当前线程的上下文能够传递到新的线程中，保持追踪上下文的一致性。
3. 结束链路追踪:
   a. 当一个操作（即一个 Span）结束时，Tracer 会标记 Span 为结束，并将其信息发送到追踪系统（如 Zipkin）。
   b. Scope 会确保上下文恢复到之前的状态，避免上下文污染。

那链路信息是如何在当前服务内以及跨服务之间传播的呢？

TracingFilter
通过观察Zipkin源码发现，在源码中有个TracingFilter实现了Filter接口。
这个过滤器的核心作用是根据 HTTP 请求的生命周期来创建、关联和结束 Span，确保链路追踪信息在请求的整个过程中得以传播，并且能够正确处理异步请求的情况。

public final class TracingFilter implements Filter {
    @Override
public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
  throws IOException, ServletException {
  HttpServletRequest req = (HttpServletRequest) request;
  HttpServletResponse res = servlet.httpServletResponse(response);

  // 检查请求是否已经存在 TraceContext
  TraceContext context = (TraceContext) request.getAttribute(TraceContext.class.getName());
  if (context != null) {
    // A forwarded request might end up on another thread, so make sure it is scoped
    Scope scope = currentTraceContext.maybeScope(context);
    try {
      chain.doFilter(request, response);
    } finally {
      scope.close();
    }
    return;
  }

  //创建新的 Span
  Span span = handler.handleReceive(new HttpServletRequestWrapper(req));

  // Add attributes for explicit access to customization or span context
  request.setAttribute(SpanCustomizer.class.getName(), span.customizer());
  request.setAttribute(TraceContext.class.getName(), span.context());
  SendHandled sendHandled = new SendHandled();
  request.setAttribute(SendHandled.class.getName(), sendHandled);

  Throwable error = null;
  
  // 创建和传递 Scope
  Scope scope = currentTraceContext.newScope(span.context());
  try {
    // any downstream code can see Tracer.currentSpan() or use Tracer.currentSpanCustomizer()
    chain.doFilter(req, res);
  } catch (Throwable e) {
    error = e;
    throw e;
  } finally {
    // When async, even if we caught an exception, we don't have the final response: defer
    if (servlet.isAsync(req)) {
      servlet.handleAsync(handler, req, res, span);
    } else if (sendHandled.compareAndSet(false, true)){
      // we have a synchronous response or error: finish the span
      HttpServerResponse responseWrapper = HttpServletResponseWrapper.create(req, res, error);
      handler.handleSend(responseWrapper, span);
    }
    scope.close();
  }
}
}

总的来说工作流程如下：

1. 判断请求是否已有 TraceContext，如果有，则使用当前上下文继续处理。
2. 如果没有 TraceContext，则创建新的 Span 并将其关联到当前线程。
3. 在过滤链中继续执行请求，确保在整个请求处理过程中链路信息得以传递。
4. 异常请求和异步请求分别在 finally 块中处理，确保在请求结束时正确地完成 Span。

为何使用@Async能正常传播链路信息？

前面提到Zipkin中的上下文信息是保存TraceContext中的，而TraceContext是基于ThreadLocal实现的。我们知道ThreadLocal中信息会在异步方法切换线程的时出现下图中上下文信息传递丢失的问题。那么在使用@Async注解的时候为什么能够正常的传播链路信息呢？
通过观察调用栈发现在调用异步方法前会执行一个TraceAsyncAspect切面。通过名字我们大概可以猜到跨线程正常传递的代码就是在这里面实现的。

TraceAsyncAspect

public class TraceAsyncAspect {
    @Around("execution (@org.springframework.scheduling.annotation.Async  * *.*(..))")
public Object traceBackgroundThread(final ProceedingJoinPoint pjp) throws Throwable {
    String spanName = name(pjp);
    Span span = this.tracer.currentSpan();
    if (span == null) {
       span = this.tracer.nextSpan();
    }
    span = span.name(spanName);
    try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span.start())) {
       span.tag(CLASS_KEY, pjp.getTarget().getClass().getSimpleName());
       span.tag(METHOD_KEY, pjp.getSignature().getName());
       return pjp.proceed();
    }
    finally {
       span.finish();
    }
}
}

通过源码我们发现它使用了 @Around 注解对标注为 @Async 的方法进行环绕处理，捕获异步方法的执行过程并跟踪它的 Span（链路追踪的基本单元）。

完整流程：

• 拦截 @Async 方法：当某个方法被 @Async 注解标注时，切面会捕获到该方法的执行。
• 获取或创建 Span：
  • 如果当前线程已有 Span（即上游调用已有追踪上下文），就直接使用它。
  • 如果当前线程没有 Span（例如，异步任务是第一次执行），就创建一个新的 Span。
• 设置 Span 的标签：根据方法的类名和方法名为 Span 添加标签，这些标签在 Zipkin UI 中会作为可视化的信息，帮助我们了解任务执行的上下文。
• 执行异步任务：通过 pjp.proceed() 执行异步方法。
• 结束 Span：无论异步方法是否成功执行，都确保在 finally 块中结束 Span，并将追踪信息发送到 Zipkin 或其他追踪系统。

其中关键点就在于tracer.withSpanInScope()方法，该方法的实现依赖于SpanInScope。

withSpanInScope 方法的作用：

上下文的推入与弹出：当你调用 withSpanInScope(span) 时，它会将 span 推入到当前线程的上下文中，确保所有后续操作都与该 Span 相关联。执行完操作后，withSpanInScope 会恢复之前的 Span 上下文，从而避免覆盖当前线程的其他 Span 上下文。

SpanInScope 内部实现：

SpanInScope 是一个实现了 AutoCloseable 接口的类，使用 try-with-resources 来确保在结束时正确清理上下文。其基本的流程如下：

• 保存当前上下文：在 withSpanInScope 被调用时，它会保存当前线程的 Span 上下文。
• 设置新的上下文：然后，新的 Span 被设置为当前线程的上下文，所有后续的操作都与该 Span 关联。
• 清理和恢复上下文：当 try-with-resources 语句块结束时，SpanInScope 会自动恢复之前的上下文，即恢复线程原先的 Span。

public class SpanInScope implements AutoCloseable {

    private final Tracer tracer;
    private final Span previousSpan;

    public SpanInScope(Tracer tracer, Span span) {
        this.tracer = tracer;
        this.previousSpan = tracer.currentSpan();  // 保存当前线程的 Span
        tracer.setSpan(span);  // 设置新的 Span
    }

    @Override
    public void close() {
        // 恢复之前的 Span
        tracer.setSpan(previousSpan);
    }
}

致此，我们就明白了为什么使用@Async还能正常传播链路信息了。

如何在使用CompletableFuture .supplyAsync()也能传播链路信息？

我们学习TraceAsyncAspect的代码后其实很容易能在CompletableFuture .supplyAsync()中模仿切面的实现方式，使得链路能够正常传播：

只需要通过Tracer获取到当前的Span，并且在线程中通过withSpanInScope()传播当前的Span即可

@Override
public void testTrace() {
    // 获取当前的Span
    Span currentSpan = tracer.currentSpan();
    CompletableFuture.supplyAsync(() -> {
        try (Tracer.SpanInScope ws = tracer.withSpanInScope(currentSpan)) {
            ResultDTO<List<AccountInfoResDTO.BaseInfo>> listResultDTO = userAccountFeignApi.listByMobiles("");
            return null;
        }
    });
}

我们知道CompletableFuture .supplyAsync()是可以指定线程池的，我们也可以使用自定义线程池，并且设置TaskDecorator在装饰器中实现我们链路信息传播。

threadPoolTaskExecutor.setTaskDecorator(new TaskDecorator() {
    Span currentSpan = tracer.currentSpan();
    @Override
    public Runnable decorate(Runnable runnable) {
        return () -> {
            try (Tracer.SpanInScope ws = tracer.withSpanInScope(currentSpan)) {
                runnable.run();
            }
        };
    }
});

一般公司会有基础框架中，会对上述完成封装，例如ZipkinHelper已经帮我们完成了跨线程传播链路信息的实现：

public Runnable wrap(Runnable runnable) {
    Span currentSpan = this.tracer.currentSpan();
    return TraceUtils.tracedRunnable(() -> {
        try (Tracer.SpanInScope scope = tracer.withSpanInScope(currentSpan)) {
            Span span = tracer.nextSpan();
            MDC.put(ZipkinKeys.TRACE_ID, span.context().traceIdString());
            MDC.put(ZipkinKeys.SPAN_ID, span.context().spanIdString());
            MDC.put(ZipkinKeys.PARENT_SPAN_ID, span.context().parentIdString());
            span.name("new_thread_started").kind(Span.Kind.SERVER)
                    .tag("thread_id", Thread.currentThread().getId() + "")
                    .tag("thread_name", Thread.currentThread().getName() + "");
            span.start();
            try (Tracer.SpanInScope ws = tracer.withSpanInScope(span)) {
                runnable.run();
            } catch (RuntimeException | Error e) {
                span.error(e);
                throw e;
            } finally {
                span.finish();
            }
        }
    });
}

因此，我们只需要在自定义线程池的装饰器中调用ZipkinHelper的wrap方法，并在使用CompletableFuture .supplyAsync()指定自定义的线程池即可：

threadPoolTaskExecutor.setTaskDecorator(new TaskDecorator() {
    @Override
    public Runnable decorate(Runnable runnable) {
        return () -> {
            try {
                zipkinHelper.wrap(runnable).run();
            }
        };
    }
});