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

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

谈谈1974

4601人浏览 · 2020-12-06 17:21:01

谈谈1974 · 2020-12-06 17:21:01 发布

1. 简介

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

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

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

在这里插入图片描述

2. 使用 Java agent 的步骤

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

创建一个指定的类作为 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)

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

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

使用 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() 方法添加了执行耗时打印，需要注意的点如下：

该实现中默认会修改静态变量 DEFAULT_METHOD 保存的指定类的指定方法的字节码，修改字节码依赖的工具为 javaassist
本实现修改方法字节码的方式是基于原来的方法 deal 复制出一个新方法，然后修改原方法名为 deal$old，最后再重新设置复制出的方法的方法体，替换为原来的方法 deal。这个过程中产生了一个新的方法 deal$old，需注意 premain 这种增强方式允许修改字节码添加新方法，agentmain 则不允许
如果存在多个 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));
    }
}