Java agent 探针技术(1)-JVM 启动时 premain 进行类加载期增强

1. 简介

在之前的文章 静态代理 一节中我们已经提到过 Java 探针技术，简单来说，在 JDK 1.5中 Java 引入了 java.lang.Instrument 包，该包提供了一些工具使得在类加载时期修改 Class 类成为了可能。这实际上就是提供了一种虚拟机级别的 AOP，其基本的原理可依据下图阐明：

在 类加载过程 一文我们提到了类加载的过程，其第一步就是加载。其实从 Java 类完整的生命周期来看，从 Java 源文件到虚拟机运行时的 Class 类，这中间还存在不少的处理过程，大致可分为如下两步。其中 Java agent 拦在 JVM 和运行时 Class 类之间，就相当于一个切面，为我们增强类功能提供了一个切入点

1. 首先 Java 文件要经过编译器编译成为 Class 字节码文件
2. Class 字节码文件经过 IO 读到 JVM 中，JVM 经过解析验证等环节最终创建出运行时的 Class 类

2. 使用 Java agent 的步骤

Java agent 的使用需要如下几个步骤：

1. 创建一个指定的类作为 Premain-Class ，类中包含 premain() 方法，该方法有如下两个声明。JVM 会优先加载方法1，加载成功忽略 2，如果1 没有，则加载 2 方法

   • public static void premain(String agentArgs, Instrumentation inst)：参数 agentArgs 是通过命令行传给 Java agent 的参数， inst 是 Java 的字节码转换工具
   • public static void premain(String agentArgs)

2. 创建 MANIFEST.MF 配置文件，将 Premain-Class指定为包含 premain() 方法的类。该配置文件通常也会将 Can-Redefine-Classes 和 Can-Retransform-Classes 配置为 true

3. 将包含 premain() 方法的类和 MANIFEST.MF 文件打包成代理 jar 包

4. 使用 java -javaagent:<jarpath>[=options] -jar xxx.jar 命令启动一个 Java 程序，并为其指定代理 jar 包

在执行第4个步骤后，目标 Java 程序启动执行 main() 方法之前，会先运行 -javaagent 参数指定的代理 jar 包内 Premain-Class 类的 premain() 方法

大部分类加载都会在main() 方法执行之后进行，这样 premain() 方法就能拦截大部分类的加载活动。没拦截到的主要是系统类，因为很多系统类必须提前加载完成，用户类的加载肯定是在premain()方法执行之后进行的

3. 使用示例

3.1 创建实现 ClassFileTransformer 接口的类

创建一个 CustomTransformer 类，该类实现了 ClassFileTransformer 接口并重写了 ClassFileTransformer#transform() 方法，主要实现的功能是为sample.ReactorMain#deal() 方法添加了执行耗时打印，需要注意的点如下：

1. 该实现中默认会修改静态变量 DEFAULT_METHOD 保存的指定类的指定方法的字节码，修改字节码依赖的工具为 javaassist
2. 本实现修改方法字节码的方式是基于原来的方法 deal 复制出一个新方法，然后修改原方法名为 deal$old，最后再重新设置复制出的方法的方法体，替换为原来的方法 deal。这个过程中产生了一个新的方法 deal$old，需注意 premain 这种增强方式允许修改字节码添加新方法，agentmain 则不允许
3. 如果存在多个 agent 修改同一个类同一个方法的字节码，需注意修改过程中产生的方法不能出现重复命名，否则会报错 Duplicate method name "deal$old" with signature "()V" in class file sample/ReactorMain

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;

/**
 * 检测方法的执行时间
 */
public class CustomTransformer implements ClassFileTransformer {

    // 被处理的方法列表
    private final static Map<String, List<String>> METHOD_MAP = new ConcurrentHashMap<>();

    private static final String DEFAULT_METHOD = "sample.ReactorMain.deal";

    private static final String CLASS_REGEX = "^(\\w+\\.)+[\\w]+$";

    private static final Pattern CLASS_PATTERN = Pattern.compile(CLASS_REGEX);

    private CustomTransformer() {
        add(DEFAULT_METHOD);
    }

    public CustomTransformer(String methodString) {
        this();
        if (!CLASS_PATTERN.matcher(methodString).matches()) {
            System.out.println("string:" + methodString + " not a method string");
            return;
        }
        add(methodString);
    }

    public void add(String methodString) {
        String className = methodString.substring(0, methodString.lastIndexOf("."));
        String methodName = methodString.substring(methodString.lastIndexOf(".") + 1);
        List<String> list = METHOD_MAP.computeIfAbsent(className, k -> new ArrayList<>());
        list.add(methodName);

    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        className = className.replace("/", ".");
        byte[] byteCode = null;
        // 判断加载的class的包路径是不是需要监控的类
        if (METHOD_MAP.containsKey(className)) {
            CtClass ctClass;
            try {
                ClassPool classPool = ClassPool.getDefault();
                // 将要修改的类的classpath加入到ClassPool中，否则可能找不到该类
                classPool.appendClassPath(new LoaderClassPath(loader));
                ctClass = ClassPool.getDefault().get(className);
                for (String methodName : METHOD_MAP.get(className)) {

                    // 得到方法实例
                    CtMethod ctMethod = ctClass.getDeclaredMethod(methodName);

                    // 创建新的方法，复制原来的方法，名字为原来的名字
                    CtMethod newMethod = CtNewMethod.copy(ctMethod, methodName, ctClass, null);
                    // 定义一个方法名用于描述修改字节码之前的原方法
                    String oldMethodName = methodName + "$old";
                    // 将原方法名称修改掉，避免和新添加的方法同名冲突
                    ctMethod.setName(oldMethodName);

                    // 构建新的方法体
                    StringBuilder bodyStr = new StringBuilder();
                    bodyStr.append("{");
                    bodyStr.append("long startTime = System.currentTimeMillis();\n");
                    // 调用原方法代码，类似于method();($$)表示所有的参数
                    bodyStr.append(oldMethodName).append("($$);\n");
                    bodyStr.append("long endTime = System.currentTimeMillis();\n");
                    String outputStr = "System.out.println(\"this method " + methodName
                            + " cost:\" +(endTime - startTime) +\"ms.\");\n";
                    bodyStr.append(outputStr);
                    bodyStr.append("}");

                    // 设置新的目标方法的方法体
                    newMethod.setBody(bodyStr.toString());
                    // 增加新方法, 原来的方法已经被修改名称为 oldMethodName，调用时会调用到新的目标方法
                    ctClass.addMethod(newMethod);
                }
                byteCode = ctClass.toBytecode();
                // ClassPool中删除该类
                ctClass.detach();
            } catch (Exception e) {
                System.out.println(e.getMessage());
                e.printStackTrace();
            }
        }
        return byteCode;
    }
}

3.2 创建使用 ClassFileTransformer 的 premain 类

创建 InstrumentMain 类，该类需要重点关注的是两个 premain() 方法。可以看到主要逻辑是在两个入参的 premain() 方法中调用 Instrumentation#addTransformer()方法，将自定义的 CustomTransformer 字节码转码器添加进去。这样在Java 程序 main() 方法执行前，每装载一个类ClassFileTransformer#transform()方法就执行一次，从而检查加载的类是否需要转换

public class InstrumentMain {
    /**
     * 该方法在main方法之前运行，与main方法运行在同一个JVM中 并被同一个System ClassLoader装载
     * 被统一的安全策略(security policy)和上下文(context)管理
     */
    public static void premain(String agentOps, Instrumentation inst) {
        System.out.println("====premain 方法执行开始");
        System.out.println(agentOps);
        inst.addTransformer(new CustomTransformer(agentOps));
        System.out.println("====premain 方法执行结束");
    }

    public static void premain(String agentOps) {
        System.out.println("====premain 方法执行开始");
        System.out.println(agentOps);
        System.out.println("====premain 方法执行结束");
    }

    public static void main(String[] args) {
    }
}

3.3 打包代理 jar 包

开发的最后一步是将包含 premain() 方法的类所在模块和 MANIFEST.MF 文件打包成代理 jar 包。IDEA 下打包 jar 包可参考博客 IDEA 打包 jar 包记录，最后创建的 MANIFEST.MF 文件内容如下，注意需要保留最后一行的空行

Manifest-Version: 1.0
Premain-Class: sample.InstrumentMain
Can-Redefine-Classes: true

3.4 测试

将如下目标类 ReactorMain 也打包成一个 jar 包，其 MANIFEST.MF 文件如下， 命名为 srcjar.jar

Manifest-Version: 1.0
Main-Class: sample.ReactorMain

public class ReactorMain {

    public static void main(String[] args) throws InterruptedException {
        deal();
    }

    public static void deal() throws InterruptedException {
        List<Integer> list = new ArrayList<>();
        list.add(1);
        list.add(2);

        int poolSize = 1;
        Long start = System.currentTimeMillis();
        CountDownLatch downLatch = new CountDownLatch(poolSize);

        Disposable disposable = Flux.range(1, 1000)
                .onBackpressureBuffer()
                .publishOn(Schedulers.elastic())
                .subscribe(null, null, downLatch::countDown);
        downLatch.await();
        disposable.dispose();

        Long end = System.currentTimeMillis();
        System.out.println("Duration:" + (end - start));
    }
}

代理 jar 包命名为 src.jar，则根据笔者 jar 包所在路径，最后的启动命令如下，可以看到修改的方法正常打印了执行耗时

java -javaagent:/Users/xxxxxx/workspace/demo/out/artifacts/src/src.jar=hello1 -jar /Users/xxxxxx/workspace/demo/out/artifacts/srcjar/srcjar.jar