概述

Javassist（Java Programming Assistant）是一个用于在运行时操作字节码的 Java 库，它允许开发人员动态生成、修改和分析 Java 类的字节码。Javassist提供了一种更高级别的 API，以 Java 代码的方式来操作字节码，而不需要直接操作复杂的字节码指令。这使得动态代码生成和修改变得更加容易和可维护。

添加maven依赖

<dependencies>
  <dependency>
    <groupId>org.javassist</groupId>
    <artifactId>javassist</artifactId>
    <version>3.28.0-GA</version>
  </dependency>

</dependencies>

新增Agent类

需要有一个premain方法，方法签名必须一样不然不认识

public class CustomerAgent {
    public static void premain(String arg,Instrumentation instrumentation){
        instrumentation.addTransformer(new ThreadLocalAdvanceTransform());
    }
}

新增字节码增强类

需要实现ClassFileTransformer 主要通过CtClass进行操作。CtClass是Javaassist操作字节码的精髓。

通过ClassPool.getDefault获取ClassPool
通过ClassPool获取、创建一个CtClass
通过CtClass操作字段、方法等
返回CtClass最新的字节数组用于生成字节码文件

public class ThreadLocalAdvanceTransform implements ClassFileTransformer {
    // 因为需要跨线程复制ThreadLocal。线程无非Runnable、Callable
    private static final String RUNNABLE = "java.lang.Runnable";
    private static final String CALLABLE = "java.util.concurrent.Callable";
    // 实现ClassFileTransformer的方法用于增强字节码文件
    @
    Override
    public byte[] transform(ClassLoader loader, String className, Class <? > classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            // 只增强指定路径下的
            if (className.startsWith("org.coderjonny") || className.startsWith("org/coderjonny")) {
                System.out.println("Clazz: " + className);
                // 将"org/coderjonny" => "org.coderjonny"转为包路径
                if (className.contains("/")) {
                    className = className.replace("/", ".");
                }
                // 获取ClassPool
                ClassPool classPool = ClassPool.getDefault();
                try {
                    // 获取需要增强的类的CtClass
                    CtClass ctClass = classPool.get(className);
                    // 获取实现的接口
                    CtClass[] interfaces = ctClass.getInterfaces();
                    String interfaceName = null;
                    for (int i = 0; i < interfaces.length; i++) {
                        CtClass anInterface = interfaces[i];
                        // 如果实现了RUNNABLE、CALLABLE就进行增强
                        if (anInterface.getName().equals(RUNNABLE)) {
                            interfaceName = RUNNABLE;
                        } else if (anInterface.getName().equals(CALLABLE)) {
                            interfaceName = CALLABLE;
                        }
                    }
                    // 如果实现了以上任意的接口
                    if (interfaceName != null) {
                        // 动态新增一个字段，用于保存当前线程的ThreadLocal
                        addField(ctClass);
                        // 在构造方法里新增逻辑储存当前线程ThreadLocal
                        advanceConstructMethod(ctClass);
                        // 在对应的run、call方法：之前新增保留新线程ThreadLocal+复制原来线程ThreadLocal到新线程；之后新增回复新线程ThreadLocal上下文
                        advanceMethod(ctClass, interfaceName);
                        // 输出最新类的字节数组
                        return ctClass.toBytecode();
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                    System.out.println("异常了");
                    throw new RuntimeException(e);
                }
            }
            return classfileBuffer;
        }
        // 新增一个成员变量保存ThreadLocal快照
    private static void addField(CtClass ctClass) throws CannotCompileException {
            CtField captured = CtField.make("private java.util.Map ThreadLocalAdvanceTransform$captured;", ctClass);
            ctClass.addField(captured);
        }
        // 在构造方法前获取当前线程的ThreadLocal
    private static void advanceConstructMethod(CtClass ctClass) throws CannotCompileException {
            CtConstructor[] declaredConstructors = ctClass.getDeclaredConstructors();
            // 构造方法
            for (int i = 0; i < declaredConstructors.length; i++) {
                CtConstructor constructor = declaredConstructors[i];
                constructor.insertBefore("ThreadLocalAdvanceTransform$captured = org.coderjonny.rpc.threadlocal.ThreadLocalManager.Transmitter.capture();");
            }

        }
        // 交换ThreadLocal
    private static void advanceMethod(CtClass ctClass, String interfaceName) throws NotFoundException,
        CannotCompileException {
            // Runnable、Callable
            CtMethod needAdvanceMethod = null;
            if (interfaceName.equals(RUNNABLE)) {
                needAdvanceMethod = ctClass.getDeclaredMethod("run");

            } else if (interfaceName.equals(CALLABLE)) {
                needAdvanceMethod = ctClass.getDeclaredMethod("call");
            }
            if (needAdvanceMethod != null) {
                // 在run、call方法之前保存新线程ThreadLocal、复制老线程ThreadLocal到新线程
                needAdvanceMethod.insertBefore("ThreadLocalAdvanceTransform$captured = org.coderjonny.rpc.threadlocal.ThreadLocalManager.Transmitter.replay(ThreadLocalAdvanceTransform$captured);");
                // 恢复新线程ThreadLocal
                needAdvanceMethod.insertAfter("org.coderjonny.rpc.threadlocal.ThreadLocalManager.Transmitter.restore(ThreadLocalAdvanceTransform$captured);", true);
            }

        }
}

打包

创建MANIFEST.MF

创建MANIFEST.MF指定agent类入口 src/main/resources/META-INF/MANIFEST.MF

Manifest-Version: 1.0
Can-Redefine-Classes: true
Can-Retransform-Classes: true
Premain-Class: org.coderjonny.rpc.agent.CustomerAgent
Boot-Class-Path: javassist-3.28.0-GA.jar

配置打包

<properties>
  <maven.configuration.manifestFile>src/main/resources/META-INF/MANIFEST.MF</maven.configuration.manifestFile>
</properties>
<build>
  <plugins>
    <plugin>
      <groupId>org.apache.maven.plugins</groupId>
      <artifactId>maven-jar-plugin</artifactId>
      <version>3.3.0</version>
      <configuration>
        <archive>
          <!-- 就是把前面的配置的MANIFEST.MF打入jar中，以指定premain的位置，否则会报：
          Failed to find Premain-Class manifest attribute in
          -->
          <manifestFile>${maven.configuration.manifestFile}</manifestFile>
        </archive>
      </configuration>
    </plugin>
  </plugins>
</build>

执行：mvn clean packag

使用

添加vm参数：

-javaagent:xxxxx/..../target/javassis-demo-1.0-SNAPSHOT.jar 
或者 
-Djavaagent=xxxxx/..../target/javassis-demo-1.0-SNAPSHOT.jar

查看增强后的字节码文件

原始代码

public class TargetClass implements Runnable{
    public TargetClass() {
    }

    @Override
    public void run() {
        System.out.println(ThreadLocalManager.getDefault().get());
    }
}

增强后

package org.coderjonny.rpc.target;

       import java.util.Map;
       import org.coderjonny.rpc.threadlocal.ThreadLocalManager;

       public class TargetClass
       implements Runnable {
           private Map ThreadLocalAdvanceTransform$captured = ThreadLocalManager.Transmitter.capture();

           @Override
           public void run() {
               try {
                   this.ThreadLocalAdvanceTransform$captured = ThreadLocalManager.Transmitter.replay(this.ThreadLocalAdvanceTransform$captured);
/*12*/             System.out.println(ThreadLocalManager.getDefault().get());
               }
               finally {
                   Object var2_3 = null;
                   ThreadLocalManager.Transmitter.restore(this.ThreadLocalAdvanceTransform$captured);
               }
           }
       }

总结

通过javassite即可实现代码无侵入性的进行功能增强。比如说：分布式链路追踪、埋点日志等等。

基于javaagent实现跨线程共享ThreadLocal

概述