基于 Javaassist 和 Java agent 实现零侵入的 CompletabelFuture 线程切换时 ThreadLocal 继承

基于 Javaassist 零侵入的 CompletabelFuture 线程切换时 ThreadLocal 继承

1 问题描述

最近在用 CompletableFuture 做性能优化时遇到一个问题，由于 CompletableFuture.supplyAsync 等方法调用时会切换线程上下文，项目中有用到了动态数据源 DynamicDataSourceContextHolder, 因此就导致异步任务无法加载到数据;

2 解决思路

这个问题的一个解决思路是创建一个工具类，对 CompletableFuture 的一些方法进行包装，比如下面的工具类: 在调用异步方法 supplyAsync 之前先获取当前线程的 dbName，然后重新生成一个 Supplier，这个 Supplier 会先设置数据源（ThreadLocal 移植），然后调用原 supplier 的 get 方法；

public class CompletableFutureUtil {
  /* 包装supplyAsync 方法 */
  public static <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier, ExecutorService executor){

    String dbName = DynamicDataSourceContextHolder._peek_();

    return CompletableFuture.supplyAsync(() -> {
      try{
        DynamicDataSourceContextHolder.push(dbName);
        return supplier.get();
      } catch (Exception e){
        throw e;
      } finally {
        DynamicDataSourceContextHolder.clear();
      }
    }, executor);
    
  }
  
  
}

这种方法的核心思想就是对 Suppiler 进行包装，可以理解为生成了一个对 Supplier 包装的包装类

public class SupplierWrapper<T> implements Supplier<T> {

  private Supplier<T> supplier;

  private String dbName;

  public SupplierWrapper(String dbName, Supplier<T> supplier){
    this.dbName = dbName;
    this.supplier = supplier;
  }

  @Override
  public T get() {
    try {
      DynamicDataSourceContextHolder.push(dbName);
      return supplier.get();
    } finally {
      DynamicDataSourceContextHolder.clear();
    }
  }
}

3 方案缺陷

最开始就是通过这种方式去解决的，但是存在一个问题，像 thenApply 这些方法都是返回一个 CompletableFuture，因此都是可以链式调用的，我个人也觉得这种链式调用比较优雅，特别是异步编排的时候，配合 stream 流写出来的代码显得很清晰；但是这种工具类，就需要在所有要执行 ComplatebleFuture 的地方都加上 CompletableFutureUtil.xxxx() 这种代码，个人觉得不是很优雅，而且无法链式调用；

然后就想通过 Spring Aop 来解决，因为本质上就是对 Supplier 的 get 方法增强，但是会发现如果对 Supplier 进行拦截，将无法使用 Lambda 表达式，而且就算可以，也无法获得 dbName，因为此时已经切换了线程了；

4 Java assist 和 Java agent 简介

基于上述分析，着重介绍本文的解决方案：Javaassist 增强方法；Java assist 可以实现在类加载到 JVM 之前对类的字节码进行修改、增强，就可以很好的解决这个问题；（这种对系统核心类库进行字节码修改的操作好像不是特别推荐，此处只是给出这种思路的具体实现，不考虑工程上的合理性）

要实现无侵入，还要使用 Java agent，先简单理解为可以在 main 方法执行前执行一些操作，实现对 CompletableFuture 类加载到 JVM 之前，先对其字节码进行修改，这样业务上使用他的一些方法时就和以前一样，实现了零侵入的方法增强；

4.1 Java assist

Java Assist是一个强大的字节码操作库，它允许开发者在运行时动态地创建、修改和分析Java字节码，而无需直接操作复杂的字节码指令。这种能力对于编写框架、AOP（面向切面编程）、动态代理、类增强等场景至关重要，特别是在那些需要对现有类进行扩展或修改，而又无法直接修改源代码的场合。例如，它常用于ORM（对象关系映射）工具中自动管理数据库访问代码，或是在应用服务器中实现热部署功能。

4.1.1CtClass与 ClassPool

在 Java Assist 中最核心的类是 CtClass，它代表一个 Java 类。而 ClassPool 则是用来管理这些 CtClass 实例的，它是创建和获取 CtClass 对象的工厂。

示例：创建一个新的类

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

public class JavaAssistExample {
    public static void main(String[] args) throws Exception {
        // 创建ClassPool实例，通常使用默认的即可
        ClassPool pool = ClassPool.getDefault();
        
        // 使用ClassPool创建一个新的类
        CtClass ctClass = pool.makeClass("com.example.MyDynamicClass");
    }
}

pool.makeClass

• 作用：此方法用于创建一个新的 CtClass 实例，即代表一个未定义的 Java 类。这个新类还没有被命名或定义任何属性和方法，是创建动态类的基础。

• 形参说明：

  • 无参数：直接调用 makeClass() 会创建一个匿名类，但通常我们会传入一个字符串参数来指定类的全限定名，如 com.example.MyClass。

• 返回值：返回一个 CtClass 对象，该对象可以进一步被编辑，如添加字段、方法、构造函数等。

4.1.2CtMethod

继续上面的例子，我们向新创建的类添加一个方法：

CtMethod method = new CtMethod(CtClass.voidType, "myMethod", new CtClass[]{}, ctClass);
method.setBody("{ System.out.println(\"Hello, World!\"); }");
ctClass.addMethod(method);

这段代码定义了一个名为 myMethod 的方法，无参数，返回类型为 void; 并在方法体中打印一条消息。

最后，可以通过 ctClass.toClass 转换为 JVM 可识别的 Class 对象，并通过反射或直接使用来实例化和调用方法。

4.1.3 CtField 和 CtConstructor

定义字段

CtField field = new CtField(pool.get("java.lang.String"), "fieldName", ctClass);
field.setModifiers(Modifier.PUBLIC); // 设置为公共访问权限_
CtConstructor constructor = ctClass.getDeclaredConstructors()[0]; // 获取默认构造函数_
constructor.insertBefore("{ this.fieldName = \"defaultValue\"; }"); // 在构造函数中初始化字段_

ctClass.addField(field);

上述代码定义了一个类型为 java.lang.String，名为 fieldName 的字段。然后通过 setModifiers 可以设置字段的访问权限，如 Modifier.PUBLIC、Modifier.PRIVATE 等。而 insertBefore 方法则允许在构造函数的开始处插入代码来初始化字段。最后把字段添加到类。

定义构造函数

有时需要修改已有构造函数的行为，或添加新的构造逻辑;

CtConstructor ctor = ctClass.getConstructors()[0]; // 获取第一个构造函数
ctor.insertBefore("{ System.out.println(\"Initializing...\"); }"); // 在构造函数开始前添加日志输出

Java Assist 通过 CtClass、CtMethod、CtField 和 CtConstructor 等类，提供了强大且灵活的 API，使开发人员能够在运行时动态地创建和修改 Java 类的结构。从定义简单的字段和方法，到复杂的构造函数逻辑调整，Java Assist 都能胜任，是进行高级代码生成和类操作的理想工具。

4.2 Java agent

Java Agent 是一种特殊的工具，它能够在 JVM 启动时或运行时动态地插入到 Java 应用程序中，无需修改应用程序源代码或编译过程。通过这种方式，Agent 可以对类文件进行转换（即字节码操作）、监控应用程序的行为，甚至修改其行为。这一切的魔法都基于 Java Instrumentation API，该 API 允许我们在类加载到 JVM 之前或之后，对其字节码进行修改。其应用场景主要有：

• 性能监控与分析：实时收集方法执行时间、内存使用情况等，帮助开发者定位性能瓶颈。
• 日志记录与跟踪：自动在方法调用前后添加日志记录，便于问题追踪和调试。
• 字节码操作与优化：如自动添加缓存逻辑、实现 AOP（面向切面编程）等功能，提高代码复用性和可维护性。
• 安全审计与加固：检查并阻止潜在的不安全代码执行，增强应用程序的安全性。
• 框架自动配置与增强：自动为应用程序集成并配置特定框架，减少手动配置工作量。

4.2.1 premain 方法

premain 是 Java Agent 的核心入口点之一，它允许我们的代码在被监测的应用程序主类的 main 方法执行之前运行。这意味着我们可以在应用程序启动的最早阶段对类进行操作，非常适合需要全局控制或初始化的操作。

要在项目中使用 premain，首先需要定义一个包含 premain 方法的类，并指定该类作为 Agent 的入口点。在 JVM 启动时，通过 -javaagent 参数指定 Agent jar 的位置及其入口类。

public class MyJavaAgent {
    public static void premain(String agentArgs, Instrumentation inst) {
        // 在这里编写你的Agent逻辑
        System.out.println("Agent loaded with arguments: " + agentArgs);
    }
}

参数解析

• agentArgs：这是传递给 Agent 的参数字符串，用于配置 Agent 的行为。
• Instrumentation inst：是 Java Instrumentation API 的核心对象，提供了修改和监控类定义的功能。

之后需要使用 maven 将打包，需要如下配置：

<plugins>
    <plugin>
        <groupId>org.apache.maven.plugins</groupId>
        <artifactId>maven-jar-plugin</artifactId>
        <version>3.3.0</version>
        <configuration>
            <archive>
                <manifest>
                    <addClasspath>true</addClasspath>
                </manifest>
                <manifestEntries>
                    <!-- 配置premain的入口 -->
                    <Premain-Class>pers.acme.agent.PreMainAgent</Premain-Class>
                    
                    <Can-Redefine-Classes>true</Can-Redefine-Classes>
                    <Can-Retransform-Classes>true</Can-Retransform-Classes>
                    <Can-Set-Native-Method-Prefix>true</Can-Set-Native-Method-Prefix>
                </manifestEntries>
            </archive>

        </configuration>
    </plugin>
</plugins>

然后使用 maven 打成 jar 包，在需要用到 java agent 的 java 程序时，加上如下参数：

-javaagent:C:\java-agent.jar=agentArgs参数

现在创建一个 demo 程序如下

public class App {
    public static void main( String[] args ) {

      System.out.println("Hello World!");
    }
}

填上以下启动参数

运行得到的结果如下：

4.2.2 agentmain 方法

如果说 premain 是 Agent 在 JVM 启动之初的先发制人，那么 agentmain 则是在应用程序运行期间灵活插入的“特工”。它允许 Agent 在 JVM 已经启动后动态地连接到正在运行的应用程序，为那些需要在运行时调整或监控的应用场景提供了可能；与 premain 类似，agentmain 也是一个静态方法，但它的使用场景更加灵活，可以在应用程序运行的任何时刻被调用。

public class DynamicAgent {
    public static void agentmain(String agentArgs, Instrumentation inst) {
        Class<?>[] cls = instrumentation.getAllLoadedClasses();
        for (Class<?> cl : cls) {
        
          if(cl.getName().equals("pers.acme.App")){
            try{
              System.out.println("Begin refactor Class[ " + cl.getName() + " ] byte code...\n");
              // 找到运行的类
              CtClass ctClass = pool.get("pers.acme.App");
              // 找到对应的方法
              CtMethod sayHello = ctClass.getDeclaredMethod("sayHello");
              // 在这个方法前后插入程序
              sayHello.insertBefore("{System.out.println(\"Agent main dynamic attached to here，before sayHello\");}");
              sayHello.insertAfter("{System.out.println(\"Agent main dynamic attached to here，after sayHello\");}");
        
              ctClass.defrost();
              // 用ctClass生成字节码重新加载per.acme.App类
              ClassDefinition classDefinition = new ClassDefinition(cl, ctClass.toBytecode());
              instrumentation.redefineClasses(classDefinition);
        
              System.out.println("Refactor Class[ " + cl.getName() + " ] byte code is done.\n");
            }catch (Exception e){
              e.printStackTrace();
            }
          }
        }
    }
}

同样需要打成 jar 包，maven 插件配置如下：

<plugins>
    <plugin>
        <groupId>org.apache.maven.plugins</groupId>
        <artifactId>maven-jar-plugin</artifactId>
        <version>3.3.0</version>
        <configuration>
            <archive>
                <manifest>
                    <addClasspath>true</addClasspath>
                </manifest>
                <manifestEntries>
                    <!-- 配置agent main入口 -->
                    <Agent-Class>pers.acme.agent.PreMainAgent</Agent-Class>
                    <Can-Redefine-Classes>true</Can-Redefine-Classes>
                    <Can-Retransform-Classes>true</Can-Retransform-Classes>
                    <Can-Set-Native-Method-Prefix>true</Can-Set-Native-Method-Prefix>
                </manifestEntries>
            </archive>

        </configuration>
    </plugin>
</plugins>

然后将其打成 jar 包，此处命名为 agent-main.jar；那么怎么使用这个 jar 呢？

被附加程序：首先因为 agent main 是运行时动态附加到正在运行的程序，那么先准备一个需要被修改的程序并运行此程序:

public class App {
    public static void main(String[] args) throws InterruptedException {
     
      for(;;) {
          sayHello();
          Thread.sleep(3000);
      }
    }
    
    public static void sayHello(){
      System.out.println("Hello Agent Main.");
    }
}

附加程序：现在可以写一个程序来动态连接到这个程序，实现方式为使用 VirtualMachine.attach(jdk1.6 之后才可以)

public class AgentMainAttachment{
    public static void main( String[] args ) {

      try {
        List<VirtualMachineDescriptor> list = VirtualMachine.list();
        for (VirtualMachineDescriptor vmd : list) {
          // 此处只附加到上述写的程序
          if (vmd.displayName().equals("pers.acme.App")) {
            VirtualMachine virtualMachine = VirtualMachine.attach(vmd.id());
            // 使用上述jar包
            virtualMachine.loadAgent("C:\agent-main.jar");
            virtualMachine.detach();
          }
        }
      } catch (Exception e) {
        e.printStackTrace();
      }
    }
 }

现在运行 AgentMainAttachment，被附加的程序 App 的运行结果如下，可以看到在运行 App 程序运行时，其 sayHello 方法被动态的改变了，其实是通过 VirtualMachine 来发送修改后的字节码，然后 JVM 重新加载 App 类的字节码。

5 零侵入 CompletableFuture中的ThreadLocal 继承

下面的 PreMainAgent 重新定义了 supplyAsync 的字节码，其中 refactorMethodByName 根据方法名重新生成对应方法，思路是先复制原有的方法，并命名为 xxx$acmeAgent,然后再重新定义方法体去调用这个 agent。（关于这里为什么不直接用 insertBefore 是因为我们在重新定义的方法体中引入的新的局部变量，会导致原有方法的局部变量表下标错乱，因此需要这么做，详细请参考 JVM 栈帧的资料）

在 refactorMethod 方法中，我们调用 generateWrapper 生成对应参数（例如 Supplier）的包装类，并将其字节码加载到 JVM。

注意：我们在重新定义方法体和和定义包装类时，用到的非 java 核心类库，全部使用 Class.forName 的方式来加载，因为 CompletableFuture 是 java 核心类库，他使用 BootstrapClassloader 加载，其不能用来加载非核心类库，因此需要用到 Class.forName 来运行时加载；

5.1 具体实现

package pers.acme.agent;

import javassist.*;

import java.lang.instrument.ClassDefinition;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class PreMainAgent {

  public static final ClassPool pool = new ClassPool();

  static {
    pool.appendSystemPath();
  }
  
  public static void premain(String agentArgs, Instrumentation instrumentation) {

    instrumentation.addTransformer(new ClassFileTransformer() {

      @Override
      public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        if (!"java/util/concurrent/CompletableFuture".equals(className)) {
          return null;
        }

        try {
          CtClass ctClass = pool.get("java.util.concurrent.CompletableFuture");

          refactorMethodByName(ctClass, "supplyAsync");
          refactorMethodByName(ctClass, "thenApply");
          refactorMethodByName(ctClass, "thenAccept");

          return ctClass.toBytecode();

        } catch (Exception e) {
          e.printStackTrace();
          throw new RuntimeException(e);
        }

      }
    }, true);

  }

  private static void refactorMethodByName(CtClass ctClass, String methodName) throws NotFoundException, CannotCompileException {
    CtMethod[] methods = ctClass.getDeclaredMethods(methodName);
    for (int i = 0; i < methods.length; i++) {
      refactorMethod(ctClass, methods[i], i);
    }
  }

  private static void refactorMethod(CtClass ctClass, CtMethod method, int offset) throws NotFoundException, CannotCompileException {
    if(Modifier.isVolatile(method.getModifiers())){
      return;
    }
    int paramCount = method.getParameterTypes().length;

    CtClass param1CtClass = method.getParameterTypes()[0];

    String param1ClzName = param1CtClass.getName();

    String param1WrapperName = toWrapperClassName(param1ClzName);

    generateWrapper(param1ClzName);

    CtMethod copy = CtNewMethod._copy_(method, ctClass, null);
    String agentMethodName = method.getName() + "$acmeAgent" + offset;
    copy.setName(agentMethodName);
    copy.setModifiers(method.getModifiers() | Modifier._FINAL _| Modifier._PRIVATE_);
    ctClass.addMethod(copy);
    System.out.println("Added method[ " + agentMethodName + " ] is done.");

    String paramsStr = IntStream._rangeClosed_(2, paramCount)
            .mapToObj(i -> "$" + i).collect(Collectors.joining(", "));

    String methodBody = "{\n" +
            " ClassLoader clzLoader = Thread.currentThread().getContextClassLoader();" +
            " Class dsClz = Class.forName(\"com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder\", true, clzLoader);\n" +
            " String dbName = (String) dsClz.getMethod(\"peek\", null).invoke(null, null);\n" +
            " Class clz = Class.forName(\"" + param1WrapperName + "\", true, clzLoader);\n" +
            " java.lang.reflect.Constructor constructor = clz.getConstructor(new Class[]{String.class," + param1ClzName + ".class});\n";
    if(paramCount > 1) {
      methodBody += " return " + agentMethodName + "((" + param1ClzName + ")constructor.newInstance(new Object[]{dbName, $1})," + paramsStr + ");\n";
    } else {
      methodBody += " return " + agentMethodName + "((" + param1ClzName + ")constructor.newInstance(new Object[]{dbName, $1}));\n";
    }

    methodBody += "}";

    System.out.println("Refactored method[ " + method.getName() + " ] body to : \n" + methodBody + "\n");

    method.setBody(methodBody);
  }

  private static CtClass generateWrapper(String interfaceName) throws NotFoundException, CannotCompileException {

    CtClass interfaceCt = pool.getCtClass(interfaceName);
    CtClass wrapperCt = null;
    try {
      wrapperCt = pool.get(toWrapperClassName(interfaceName));
    } catch (javassist.NotFoundException e){

    }

    if(wrapperCt == null) {
      wrapperCt = pool.makeClass(toWrapperClassName(interfaceName));
      wrapperCt.setGenericSignature(interfaceCt.getGenericSignature());
      wrapperCt.addInterface(interfaceCt);

      // 添加一个私有字段（域）
      CtField dbNameField = new CtField(pool.getCtClass("java.lang.String"), "dbName", wrapperCt);
      dbNameField.setModifiers(Modifier.PRIVATE); // 设置为私有
      wrapperCt.addField(dbNameField); // 将字段添加到类中

      CtField funcField = new CtField(pool.getCtClass(interfaceName), "func", wrapperCt);
      funcField.setModifiers(Modifier.PRIVATE); // 设置为私有
      wrapperCt.addField(funcField); // 将字段添加到类中


      // 定义构造函数的参数类型
      CtClass[] constructorParameterTypes = new CtClass[]{pool.get("java.lang.String"), interfaceCt};
      // 创建构造函数
      CtConstructor wrapperConstructor = new CtConstructor(constructorParameterTypes, wrapperCt);
      wrapperConstructor.setModifiers(Modifier._PUBLIC_);
      // 设置构造函数的主体
      wrapperConstructor.setBody("{\n" +
              " this.dbName = $1;\n" +
              " this.func = $2;\n" +
              "}");
      wrapperCt.addConstructor(wrapperConstructor);

      CtMethod[] declaredMethods = interfaceCt.getDeclaredMethods();
      for (CtMethod method : declaredMethods) {
        if(!Modifier.isAbstract(method.getModifiers())){
          continue;
        }
        CtMethod wrapperMethod = new CtMethod(method.getReturnType(), method.getName(), method.getParameterTypes(), wrapperCt);
        String methodBody = "{\n" +
                " ClassLoader clzLoader = Thread.currentThread().getContextClassLoader();\n" +
                " Class dsClz = null;\n" +
                " try {\n" +
                "   dsClz = Class.forName(\"com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder\", true, clzLoader);\n" +
                "   dsClz.getMethod(\"push\", new Class[]{String.class}).invoke(null, new Object[]{this.dbName});\n";

        if (!method.getReturnType().getName().equals("void")) {
          methodBody += "   return this.func." + method.getName() + "($$);\n";
        } else {
          methodBody += "   this.func." + method.getName() + "($$);\n";
        }

        methodBody += " } catch(Exception e) {\n" +
                "   e.printStackTrace();\n" +
                "   throw e;" +
                " } finally {\n" +
                "   dsClz.getMethod(\"clear\", null).invoke(null, null);\n" +
                " }\n";
        methodBody += "}";
        wrapperMethod.setBody(methodBody);

        wrapperCt.addMethod(wrapperMethod);

      }
    }

    try {
      wrapperCt.toClass();
    } catch (CannotCompileException e){
      // 被重复加载
    }
    return wrapperCt;
  }

  private static String toWrapperClassName(String interfaceName){
    String[] names = interfaceName.split("\\.");
    String wrapperName = "pers.acme.agent.wrapper." + names[names.length - 1] + "Wrapper";
    return wrapperName;
  }
}

将此程序打成 jar 包，命名为 completable-future-refactor.jar

现在运行一个 CompletableFuture 的程序:

package pers.acme;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class App {
  public static void main(String[] args) throws InterruptedException {
    ExecutorService executorService = Executors.newFixedThreadPool(1);


    DynamicDataSourceContextHolder.push("acme_database");

    CompletableFuture<String> helloFuture = CompletableFuture.supplyAsync(() -> {

      System.out.println("DynamicDataSourceContextHolder.peek() = " + DynamicDataSourceContextHolder._peek_());

      return "Hello";
    }, executorService);

    System.out.println(helloFuture.join());

    executorService.shutdown();
  }

}

先不用 javaagent，结果如下：

然后在启动参数上加上-javaagent:C:\completable-future-refactor.jar，运行结果如下:

可以看到 CompletableFuture 的字节码已经被改变。