How to make SOFATracer support @Async in Spring

前言

最近公共团队的SkyWalking不咋灵光，也没人维护了，就寻思的捣鼓一套简单的能自用的，用于请求追踪和Span耗时分析的组件，能接受一定的代码侵入。调研了一番以后发现SOFATracer还是挺满足我的需要的，虽然感觉现在他们更新的也不勤快了，也有些比较奇妙的bug，但好在都比较好解决。

这类APM工具一般都是用ThreadLocal来存储上下文，由于ThreadLocal的特性，在异步场景下一般都需要做线程的包装才能正常使用，SOFATracer提供了一些线程池的包装，但对于Spring-Context下的@Async的支持，还是也需要做一些代码上的改动

（对比下来，确实SkyWalking还是最易用的，也没有侵入性，也不需要考虑各种问题）

Async的配置

在Spring项目中要能使用@Async实现异步，一般需要进行以下几步

在配置类里加入@EnableAsync的注解，实现对AsyncConfigurationSelector的导入
配置一个默认线程池用于执行@Async的任务，可以通过实现org.springframework.scheduling.annotation.AsyncConfigurer的接口，或是在IoC容器里声明一个org.springframework.core.task.TaskExecutor的类的Bean（有且只有一个），或是在IoC容器里声明一个名字为taskExecutor的java.util.concurrent.Executor类的Bean

关于默认线程池的获取逻辑，可以查看org.springframework.aop.interceptor.AsyncExecutionAspectSupport#getDefaultExecutor中的实现

对于SpringBoot的项目来说，SpringBoot自动装配时会装配一个ThreadPoolTaskExecutor，具体装配类代码在org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration中

实现方法

知道了@Async线程池的来源，要怎么实现也就简单了，对执行的线程任务做包装即可。包装有两种思路，一种思路是使用对线程池进行包装，使其在提交任务时对执行的任务进行包装；第二种思路是对执行的任务直接进行包装

由于SOFATracer本身没有对java.util.concurrent.Executor线程池包装的实现，需要我们自己先实现一个

/**
 * @author xYohn
 * @date 2024/4/10
 */
public class SofaWrapExecutor implements Executor {

    private final Executor delegate;

    public SofaWrapExecutor(Executor delegate) {
        this.delegate = delegate;
    }

    /**
     * @param command the runnable task
     */
    @Override
    public void execute(Runnable command) {
        delegate.execute(new SofaTracerRunnable(command));
    }
}

对于使用org.springframework.scheduling.annotation.AsyncConfigurer方式配置默认线程池的实现方式

org.springframework.scheduling.annotation.AsyncConfigurer#getAsyncExecutor要求的返回类型是Executor，所以我们只需要在原先的线程池配置上，return一个被上述类SofaWrapExecutor包装的线程池即可

对于在IoC容器里声明一个org.springframework.core.task.TaskExecutor的类的Bean，同样可以参考上述方法，只需把上述代码中的Executor统一替换成TaskExecutor即可
对于指定自定义的线程池执行的方式(如在加上@Async注解时指定了其他线程池)(@Async("pool"))，实际上跟1的情况一样，需要在这个Bean实例化的逻辑处加上上述包装即可

对于SpringBoot的项目且使用SpringBoot自动装配的线程池的情况，SpringBoot自动装配的ThreadPoolTaskExecutor线程池提供了org.springframework.core.task.TaskDecorator用于对线程进行包装处理。同时，SpringBoot提供了org.springframework.boot.task.TaskExecutorBuilder用于自定义该线程池

其中org.springframework.core.task.TaskDecorator通过注入的形式用于org.springframework.boot.task.TaskExecutorBuilder的构建，因此，可以通过构造org.springframework.core.task.TaskDecorator的实例，使用SofaTracerRunnable对线程的包装，然后注入到IoC容器中实现该效果，示例代码如下

/**
 * @author xYohn
 * @date 2024/4/10
 */
@Configuration
public class TaskDecoratorConfiguration {

    @Bean
    public TaskDecorator taskDecorator(){
        return SofaTracerRunnable::new;
    }
}

至此，SOFATracer对@Async的支持就完成了。