如何实现一个简单高效的链路追踪系统

931 阅读6分钟

对于大型服务端应用来说,链路追踪的功能是必不可少的。我们可以通过一个TraceId来定位整个请求的处理链路,从而来帮助我们快速定位问题。

市面上是有一些链路追踪框架的,比如Skywalking,但是了解下来发现侵入性挺强的,而且各种配置真的繁琐,我也不需要那么多乱七八糟的功能。所以我自己写了套简单但是完全满足链路追踪的能力。

核心功能

当你对小爱同学说了一句话”今天天气“,如果没有收到响应,服务端同学就要拿你当前的requestId去查找问题。这个requestId是客户端生成的,然后在内部和外部流转,小爱的处理流程非常复杂,中间串行并行的服务能有上百个。如果没有这个requestId,调查问题难于登天。但是有了这个requestId后,服务端的人就可以拿这个requestId直接去Kibana上面直接去搜索日志了。

所以我的诉求就是

  1. 生成traceId:我们可以直接使用uuid即可
  2. 在日志中显示traceId
  3. traceId写到es中去:当然如果接入了elk,由于日志被收集了,那么traceId也自然被写入了。
  4. 在各个服务间传递traceId

在日志中显示traceId

接入流程

在日志中显示traceId,主要做两件事

一、设置日志格式: %X{key}

<property name="pattern"
              value="%d{yyyy-MM-dd HH:mm:ss.SSS} [%thread] %logger{50} %X{key} [%.-5level]: %msg%n"/>

最好打印出来带traceId的日志如下

2024-05-13 10:18:34.207 [ConsumeMessageThread_resource_usage_stat_consumer_1] c.i.c.g.c.s.s.impl.ResourceUsageStatServiceImpl  ed79ce56d87547cf82dbe268a2b8987b [INFO]:

二、将traceId绑定到线程

其实就是调用MDC.put(key, traceId)即可,比如对接客户端的一般是个HTTP服务,那么我们添加个拦截器即可。

假设你用的springboot,那么可以添加下面的拦截器

public class TraceIdInterceptor implements HandlerInterceptor {
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response,
                            Object handler) {
        String traceId = request.getHeader(TraceIdUtils.TRACE_KEY);
        if (StringUtils.isEmpty(traceId)) {
            TraceIdUtils.genAndSetTraceId();
        } else {
            TraceIdUtils.setTraceId(traceId);
        }
        return true;
    }

    @Override
    public void postHandle(HttpServletRequest request,HttpServletResponse response,
                           Object handler, ModelAndView modelAndView) {
        TraceIdUtils.clear();
    }
}
import org.apache.commons.lang3.StringUtils;
import org.apache.skywalking.apm.toolkit.trace.TraceContext;
import org.slf4j.MDC;

import java.util.Optional;
import java.util.UUID;

public class TraceIdUtils {

    public static final String TRACE_KEY = "trace_id";

    public static String generateTraceId() {
        return UUID.randomUUID().toString().replaceAll("-", "");
    }

    public static void genAndSetTraceId() {
        setTraceId(generateTraceId());
    }

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

    public static void setTraceId(String traceId) {
        Optional.ofNullable(traceId).ifPresent(x -> MDC.put(TRACE_KEY, traceId));
    }

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

    public static void remove() {
        MDC.remove(TRACE_KEY);
    }
}

MDC介绍

MDC是SLF4J提供的一个上下文信息存储工具,可以在日志输出中添加额外的上下文信息。MDC内部持有个ThreadLocal<Map<String,String>> copyOnThreadLocal。下面是它的put和get方法

public void put(String key, String val) throws IllegalArgumentException {
    if (key == null) {
        throw new IllegalArgumentException("key cannot be null");
    }

    Map<String, String> oldMap = copyOnThreadLocal.get();
    Integer lastOp = getAndSetLastOperation(WRITE_OPERATION);

    if (wasLastOpReadOrNull(lastOp) || oldMap == null) {
        Map<String, String> newMap = duplicateAndInsertNewMap(oldMap);
        newMap.put(key, val);
    } else {
        oldMap.put(key, val);
    }
}
private Map<String, String> duplicateAndInsertNewMap(Map<String, String> oldMap) {
    Map<String, String> newMap = Collections.synchronizedMap(new HashMap<String, String>());
    if (oldMap != null) {
        // we don't want the parent thread modifying oldMap while we are
        // iterating over it
        synchronized (oldMap) {
            newMap.putAll(oldMap);
        }
    }

    copyOnThreadLocal.set(newMap);
    return newMap;
}

public String get(String key) {
    final Map<String, String> map = copyOnThreadLocal.get();
    if ((map != null) && (key != null)) {
        return map.get(key);
    } else {
        return null;
    }
}

​ 我们可以通过调用put(key, traceId)来讲traceId绑定到当前线程,然后在打印日志的时候日志框架识别到了占位符%X{key}就调用get(key)方法来取出当前线程绑定的traceId打印出来。

服务间传递traceId

这里我简单介绍下在各个服务通信中如何传递traceId。比如传递给Http服务,传递给rpc服务,传递给mq以及异步处理时如何在线程之间传递。

传递给Http服务

前面我介绍在http服务作为接收端如何处理traceId,简单来说就是判断当前请求头中是否有traceId,如果有责设置到mdc中,如果没有就生成后再设置到mdc中。

所以作为发送端,则需要将traceId写入到请求头中,比如我们常用OkHttp来进行http调用,那么我们可以自定义一个拦截器,在拦截器中中设置请求头

private static final List<Interceptor> DEFAULT_INTERCEPTORS =
            Lists.newArrayList(new OkhttpTraceInterceptor());
private static final ConnectionPool DEFAULT_CONNECTION_POOL =
            new ConnectionPool(25, 5, TimeUnit.MINUTES);

private static OkHttpClient client;
static {
    HttpClientConfig config = new HttpClientConfig()
            .setConnectTimeout(Duration.ofSeconds(30))
            .setReadTimeout(Duration.ofMinutes(1))
            .setWriteTimeout(Duration.ofMinutes(1))
            .setConnectionPool(DEFAULT_CONNECTION_POOL)
            .setInterceptors(DEFAULT_INTERCEPTORS);
    OkHttpClient.Builder builder = new OkHttpClient().newBuilder()
                .connectTimeout(config.getConnectTimeout())
                .readTimeout(config.getReadTimeout())
                .writeTimeout(config.getWriteTimeout())
                .connectionPool(config.getConnectionPool());
    config.getInterceptors().forEach(builder::addInterceptor);
    client = builder.build();
}

public class OkhttpTraceInterceptor implements Interceptor {

    @Override
    @Nonnull
    public Response intercept(Chain chain) throws IOException {
        Request.Builder builder = chain.request().newBuilder();
        String traceId = TraceIdUtils.getTraceId();
        if (StringUtils.isNotEmpty(traceId)) {
            builder.addHeader(TraceIdUtils.TRACE_KEY, traceId);
        }
        Request request = builder.build();
        return chain.proceed(request);
    }
}

传递给RocketMQ

在消息队列之间的传递这里我选择的是rocketmq,其他mq也类似。

发送端

一、定义一个钩子,在每次发送消息之前从当前线程获取traceId设置到消息的用户属性当中即可。

public class TraceSendMsgHook implements SendMessageHook {
    @Override
    public String hookName() {
        return "TraceSendMsgHook";
    }

    @Override
    public void sendMessageBefore(SendMessageContext context) {
        String traceId = TraceIdUtils.getTraceId();
        if (StringUtils.isNotEmpty(traceId)) {
            context.getMessage().putUserProperty(TraceIdUtils.TRACE_KEY, traceId);
        }
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {

    }
}

二、构建生产者的时候将这个钩子注册进去

DefaultMQProducer producer = new DefaultMQProducer(this.producerName);
producer.getDefaultMQProducerImpl().registerSendMessageHook(new TraceSendMsgHook());

需要注意的是这种方式只针对同步发送才起作用。

如果使用异步发送消息,那么mq会使用内部的线程池来发送消息,此时发送端就不能从当前线程获取traceId了。这种情况,可以自定义发送端的线程池(自己写的带链路追踪的线程池)。

接收端

一、定义一个消费者钩子,在每次接收消息的时候,从用户属性中取traceId,并绑定到当前线程中去

public class TraceConsumeMessageHook implements ConsumeMessageHook {
    @Override
    public String hookName() {
        return "TraceConsumeMessageHook";
    }

    @Override
    public void consumeMessageBefore(ConsumeMessageContext context) {
        context.getMsgList().stream().findFirst()
                .ifPresent(x -> TraceIdUtils.setTraceId(x.getUserProperty(TraceIdUtils.TRACE_KEY),
                        x.getUserProperty(TraceIdUtils.USER_KEY)));
    }

    @Override
    public void consumeMessageAfter(ConsumeMessageContext context) {
        TraceIdUtils.clear();
    }
}

二、构建消费者的时候,注册这个钩子即可

DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(consumerGroup);
consumer.getDefaultMQPushConsumerImpl().registerConsumeMessageHook(new TraceConsumeMessageHook());

异步处理

在代码里面如果使用线程池那么就会导致我们打印不出来traceId了,因为这时候MDC获取的当前线程是线程池的线程而不是调用线程池的线程(主线程),而traceId是在主线程中的,这种情况如何处理呢?

很简单,将主线程中的traceId拷贝到线程池的执行线程中即可,比如我们常使用下面的代码进行异步处理

executorService.submit(Callable);
executorService.execute(Runnable)

executorService是我们的线程池,我们只要重新定义下callable和runnable,在创建的时候就将主线程的traceId拷贝保存起来,当线程池中的线程执行任务的时候将traceId再设置给当前线程即可。

public class TraceLambdaHook {
    protected final Map<String, String> context;

    // 创建任务的时候拷贝主线程的上下文信息,内部实现是
    // Map<String, String> hashMap = copyOnThreadLocal.get();
    //     if (hashMap == null) {
    //         return null;
    //     } else {
    //         return new HashMap<String, String>(hashMap);
    //     }
    protected TraceLambdaHook() {
        this.context = MDC.getCopyOfContextMap();
    }
    protected void setContext() {
        if (MapUtils.isNotEmpty(context)) {
            MDC.setContextMap(context);
        }
    }

    protected void clearContext() {
        if (MapUtils.isNotEmpty(context)) {
            MDC.clear();
        }
    }
}

public class TraceCallable<T> extends TraceLambdaHook implements Callable<T> {
    private final Callable<T> delegate;

    public TraceCallable(Callable<T> delegate) {
        this.delegate = delegate;
    }

    @Override
    public T call() throws Exception {
        try {
            // 执行任务的时候将上下文信息设置当当前线程(线程池的线程)
            setContext();
            return delegate.call();
        } finally {
            // 执行完了清理上下文信息
            clearContext();
        }
    }
}

public class TraceRunnable extends TraceLambdaHook implements Runnable {
    private final Runnable delegate;
    public TraceRunnable(Runnable delegate) {
        this.delegate = delegate;
    }

    @Override
    public void run() {
        try {
            setContext();
            delegate.run();
        } finally {
            clearContext();
        }
    }
}

当然你也可以参考上面的方式来写TraceSupplier等函数式接口。

如果你要使用java8提供的CompletableFuture::supplyAsync

public static <U> CompletableFuture<U> supplyAsync(Supplier<U> supplier,
                                                       Executor executor) 

那么只需要继承ThreadPoolExecutor,重写下execute()和submit()方法

public class TraceThreadPoolExecutor extends ThreadPoolExecutor {

    public TraceThreadPoolExecutor(
            int corePoolSize, int maximumPoolSize, long keepAliveTime,
            @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public TraceThreadPoolExecutor(
            int corePoolSize, int maximumPoolSize, long keepAliveTime,
            @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    @Override
    public void execute(@Nonnull Runnable task) {
        super.execute(new TraceRunnable(task));
    }

    @Override
    public <T> Future<T> submit(@Nonnull Callable<T> task) {
        return super.submit(new TraceCallable<>(task));
    }

}

或者你不需要重写线程池,使用自定义的函数式接口即可。