【JavaAgent】字节码编程 - 方法运行时间实现

使用JavaAgent实现一个简单的需求：方法的运行时间。这是一个很简单的例子，我们的目的是了解各种库之间的差异，做出正确的选择。

我们将学习：

• 使用ASM, Javassist和byte-buddy库实现
• 编写测试用例，运行查看效果
• 设置环境变量切换实现

使用版本

• JDK 11
• asm 9.2
• javassist 3.28
• byte-buddy 1.12.3

我们要确定目标类，也就是哪些类需要处理。本项目确定一个目标类 TargetClass.java, 代码如下：

@Slf4j
public class TargetClass {

    public void method1() throws InterruptedException{
        log.info(">>> method1 called ");
        Thread.sleep((new Random()).nextInt(1000));
    }

    public String method2(){
        log.info(">>> method2 called ");
        
        try {
            Thread.sleep((new Random()).nextInt(1000));
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        return "end";
    }
}

我们要了解这个类的方法运行时间，这里需要注意: method1 没有返回值。字节码编程时，这两个方法的处理有所不同。

我们还需要程序入口，并且创建了 TargetClass 实例，运行method1和method2方法，Main.java代码如下：

@Slf4j
public class Main {
    public static void main(String[] args) throws InterruptedException {
        log.info(">>> Main is running -> {}", Main.class.getName());

        TargetClass target = new TargetClass();
        target.method1();
        target.method2();
    }
}

运行程序，可以看见目标方法的执行情况。

我们需要一个功能统计运行时间，秒表（StopWatch.java）实现如下：

@Slf4j
public class StopWatch {
    
    public static class StaticClazz {
        static ThreadLocal t = new ThreadLocal();
    
        public static void start() {
            t.set(System.currentTimeMillis());
        }
    
        public static void end() {
            final long elapseOfTime = System.currentTimeMillis() - t.get();
            log.info("{} elapse of time: {}", Thread.currentThread().getStackTrace()[2] , elapseOfTime);
    
            t.remove();          
        }
    
    }

    
    public static class Clazz {
        private final String methodName;
        private long start;
    
        public Clazz(String methodName) {
            this.methodName = methodName;
        }

        public  void start() {
            start = System.currentTimeMillis();
        }
    
        public void end() {
            final long elapseOfTime = System.currentTimeMillis() - start;
            log.info("{} elapse of time: {}", methodName , elapseOfTime);
       
        }
    
    }

}

这里有两种实现：类和静态类，为什么呢？字节码编程是很专业的，需要学习库的语法，使用静态类可以降低实现需求的难度。

我们先实现需求，对目标类方法编程，ASM库实现（TransformerWithASM.java）:

public class TransformerWithASM implements Transformer{
    private final static String TARGET_CLASS = "io/github/kavahub/learnjava/TargetClass";

    @Override
    public void transform(String args, Instrumentation inst) {
        inst.addTransformer(new ClassFileTransformer() {

            @Override
            public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined,
                    ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                log.info("transform class - " + className);

                // 仅处理TargetClass类
                if (className.equals(TARGET_CLASS)) {
                    ElapseOfTimeClassWriter writer = new ElapseOfTimeClassWriter(classfileBuffer);
                    return writer.perform();
                }

                return ClassFileTransformer.super.transform(loader, className, classBeingRedefined, protectionDomain,
                        classfileBuffer);
            }

        });
    }

    public static class ElapseOfTimeClassWriter {
        private final ClassReader reader;
        private final ClassWriter writer;

        public ElapseOfTimeClassWriter(byte[] contents) {
            reader = new ClassReader(contents);
            writer = new ClassWriter(reader, 0);
        }

        public byte[] perform() {
            ElapseClassAdapter elapseClassAdapter = new ElapseClassAdapter(writer);
            reader.accept(elapseClassAdapter, 0);
            return writer.toByteArray();
        }
    }

    public static class ElapseClassAdapter extends ClassVisitor {
        public ElapseClassAdapter(ClassVisitor classVisitor) {
            super(ASM7, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature,
                String[] exceptions) {
            MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);

            return new ElapseMethodAdapter(methodVisitor);

        }

    }

    public static class ElapseMethodAdapter extends MethodVisitor {
        private final static String owner = "io/github/kavahub/learnjava/StopWatch$StaticClazz";

        public ElapseMethodAdapter(MethodVisitor methodVisitor) {
            super(ASM7, methodVisitor);
        }

        
        @Override
        public void visitCode() {
            // 方法开始时，插入StopWatch代码，调用start方法
            mv.visitMethodInsn(INVOKESTATIC, owner, "start", "()V", false);
            super.visitCode();
        }

        @Override
        public void visitInsn(int opcode) {
            if ((opcode >= IRETURN && opcode <= RETURN)) {
                // 方法返回时, 插入StopWatch代码，调用end方法
                visitMethodInsn(INVOKESTATIC, owner, "end", "()V", false);
            }
            mv.visitInsn(opcode);
        }

    }

}

为了降低实现难度，我们使用了秒表（StopWatch）的静态类。

Javassist实现

public class TransformerWithJavassist  implements Transformer {
    private final static String TARGET_CLASS = "io/github/kavahub/learnjava/TargetClass";

    @Override
    public void transform(String args, Instrumentation inst) {
        inst.addTransformer(new ClassFileTransformer() {

            @Override
            public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined,
                    ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                log.info("transform class - {}" + className);

                // 仅处理TargetClass类
                if (className.equals(TARGET_CLASS)) {
                    ElapseOfTimeClassWriter writer = new ElapseOfTimeClassWriter(className);
                    return writer.write();
                }

                return ClassFileTransformer.super.transform(loader, className, classBeingRedefined, protectionDomain,
                        classfileBuffer);
            }

        });
    }

    public static class ElapseOfTimeClassWriter {
        private final static String STOPWATCH_START = "n final io.github.kavahub.learnjava.StopWatch.Clazz stopwatch = new io.github.kavahub.learnjava.StopWatch.Clazz("%s");n"
                + "n stopwatch.start();n";
        private final static String STOPWATCH_END = "n stopwatch.end();n";

        private final String className;

        public ElapseOfTimeClassWriter(String className) {
            this.className = className.replace("/", ".");
        }

        public byte[] write() {
            try {
                return wirte0();
            } catch (NotFoundException | CannotCompileException | IOException e) {
                throw new RuntimeException(e);
            }
        }

        private byte[] wirte0() throws NotFoundException, CannotCompileException, IOException {
            CtClass ctclass = ClassPool.getDefault().get(className);
            for (CtMethod ctMethod : ctclass.getDeclaredMethods()) {
                CtClass returnType = ctMethod.getReturnType();
                // 无返回值方法
                if (CtClass.voidType.equals(returnType)) {
                    String methodName = ctMethod.getName();
                    // 新定义一个方法
                    String newMethodName = "elapse$" + methodName;
                    // 将原来的方法名字修改
                    ctMethod.setName(newMethodName);
                    // 创建新的方法，复制原来的方法，名字为原来的名字
                    CtMethod newMethod = CtNewMethod.copy(ctMethod, methodName, ctclass, null);
                    // 构建新的方法体
                    StringBuilder bodyStr = new StringBuilder();
                    bodyStr.append("{");
                    bodyStr.append(String.format(STOPWATCH_START, methodName));
                    // 调用原有代码，类似于method();($$)表示所有的参数
                    bodyStr.append(newMethodName + "($$);n");
                    bodyStr.append(STOPWATCH_END);
                    bodyStr.append("}");

                    // 替换新方法
                    newMethod.setBody(bodyStr.toString());
                    // 增加新方法
                    ctclass.addMethod(newMethod);
                }
            }
            // 修改后的方法列表 会发现多了一个方法
            log.info(" after update method list ...");
            for (CtMethod ctMethod : ctclass.getDeclaredMethods()) {
                log.info("Method name - {}", ctMethod.getName());
            }
            return ctclass.toBytecode();
        }
    }

}

byte-buddy 实现

public class TransformerWithByteBuddy implements Transformer {
    private final static String TARGET_CLASS = "io.github.kavahub.learnjava.TargetClass";

    private AgentBuilder.Transformer transformer = (builder, typeDescription, classLoader, javaModule) -> {
        return builder
                // 拦截任意方法
                .method(ElementMatchers.any())
                // 委托
                .intercept(MethodDelegation.to(ElapseOfTimeWriter.class));
    };

    private AgentBuilder.Listener listener = new AgentBuilder.Listener() {

        @Override
        public void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
            log.info("onDiscovery - {}", typeName);
        }

        @Override
        public void onTransformation(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module,
                boolean loaded, DynamicType dynamicType) {
            log.info("onTransformation - {}", typeDescription);
        }

        @Override
        public void onIgnored(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module,
                boolean loaded) {
            log.info("onIgnored - {}", typeDescription);
        }

        @Override
        public void onError(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded,
                Throwable throwable) {
            log.info("onError - {}", typeName);
        }

        @Override
        public void onComplete(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
            log.info("onComplete - {}", typeName);
        }
    };

    @Override
    public void transform(String args, Instrumentation inst) {
        new AgentBuilder.Default()
                // 指定需要拦截的类
                .type(ElementMatchers.named(TARGET_CLASS)) 
                .transform(transformer)
                .with(listener)
                .installOn(inst);

    }

    public static class ElapseOfTimeWriter {
        @RuntimeType
        public static Object wirte(@Origin Method method, @SuperCall Callable callable) throws Exception {
            final StopWatch.Clazz stopWatch = new StopWatch.Clazz(method.getName());
            stopWatch.start();
            try {
                // 原有函数执行
                return callable.call();
            } finally {
                stopWatch.end();
            }
        }
    }

}

 三个实现类中，都有ElapseOfTimeWriter静态类，这个类的就是字节码编程，实现需求的。

我们有三种实现，用户可以选择其中的任意一种。如果实现这个功能，需要实现Provider（提供者），使用环境变量决定使用哪种实现，TransformerProvider.java 代码如下：

public class TransformerProvider implements Supplier {
    public final static String CLASS_FILE_TRANSFORMER_KEY = "transformer_class";
    private final String CLASS_FILE_TRANSFORMER_VALUE;

    public final static TransformerProvider INSTANCE = new TransformerProvider();

    private TransformerProvider() {
        CLASS_FILE_TRANSFORMER_VALUE = System.getenv(CLASS_FILE_TRANSFORMER_KEY);
    }

    @Override
    public Transformer get() {
        String className = CLASS_FILE_TRANSFORMER_VALUE;
        if (className == null) {
            // 默认
            className = TransformerWithASM.class.getName();
        }

        // 反射创建
        Class clazz = null;
        try {
            clazz = Class.forName(className);
            if (Transformer.class.isAssignableFrom(clazz)) {
                return (Transformer)clazz.getDeclaredConstructor().newInstance();
            }
            log.error("Is not a correct subclass of ClassFileTransformer -> {}", className);
        } catch (ClassNotFoundException | InstantiationException | IllegalAccessException 
            | IllegalArgumentException | InvocationTargetException | NoSuchMethodException | SecurityException e) {
            log.error("Fail to create object", e);
        }

        

        return null;
    }

}

需求已经完毕执行完成。现在我们要与JavaAgent集成了，PremainAgent.java代码如下：

@Slf4j
public class PremainAgent {
    public static void premain(String args, Instrumentation inst){
        log.info("Agent called - {}", PremainAgent.class.getName());
        
        Transformer transformer = TransformerProvider.INSTANCE.get();
        if (transformer != null) {
            transformer.transform(args, inst);
            log.info("Transformer registered successfully -> {}", transformer.getClass().getName());
        } else {
            log.warn("Agent failure, because of transformer is not configured correctly");
        }
        
    }
}

集成完成了。我们使用maven-jar-plugin插件打包项目，POM配置如下：

                org.apache.maven.plugins
                maven-jar-plugin
                
                    
                        
                            true
                        
                        
                            
                                io.github.kavahub.learnjava.PremainAgent
                            

                            true
                            true
                        
                    
                
            

运行Maven命令打包：

mvn clean install

在 target 目录下，生成 elapse-of-time.jar

通常使用命令行，格式如下：

java -javaagent:path/to/elapse-of-time.jar -cp %classpath% io.github.kavahub.learnjava.Main

上面只是格式，不是最终的命令。其中classpath最繁琐。

我们使用测试用例的方式：在测试种运行命令行。需要增加Main.java的功能：

@Slf4j
public class Main {
    
    ......

    
    public static ProcessBuilder build() throws IOException, InterruptedException {
        ProcessBuilder builder = new ProcessBuilder();
        addJavaBin(builder);
        // 注意：javaagent要放在前面
        addJavaAgent(builder);
        addClasspath(builder);
        addClassMain(builder);
    
        builder.inheritIO();
        return builder;
      }

    private static void addClassMain(ProcessBuilder builder) {
        String className = Main.class.getCanonicalName();
        builder.command().add(className);
    }

    private static void addClasspath(ProcessBuilder builder) {
        String classpath = System.getProperty("java.class.path");
        builder.command().add("-cp");
        builder.command().add(classpath);
    }

    private static void addJavaAgent(ProcessBuilder builder) {
        Path javaagent = Paths.get("target", "elapse-of-time.jar");
        builder.command().add("-javaagent:" + javaagent.toAbsolutePath().toString());
    }

    private static void addJavaBin(ProcessBuilder builder) {
        String javaHome = System.getProperty("java.home");
        String javaBin = javaHome + File.separator + "bin" + File.separator + "java";
        builder.command().add(javaBin);
    }
}

测试用例 TargetClassManualTest.java 代码如下：

public class TargetClassManualTest {

    @Test
    public void givenDefault_whenRunProgream_thenCheckConsole() throws IOException, InterruptedException {
        // 启动程序
        ProcessBuilder builder = Main.build();
        Process process = builder.start();

        // 等待程序运行完成
        while(process.isAlive()) {}
    }

    @Test
    public void givenJavassist_whenRunProgream_thenCheckConsole() throws IOException, InterruptedException {
        // 启动程序
        ProcessBuilder builder = Main.build();
        builder.environment().put("transformer_class", "io.github.kavahub.learnjava.TransformerWithJavassist");
        Process process = builder.start();

        // 等待程序运行完成
        while(process.isAlive()) {}
    }
    

    @Test
    public void givenASM_whenRunProgream_thenCheckConsole() throws IOException, InterruptedException {
        // 启动程序
        ProcessBuilder builder = Main.build();
        builder.environment().put("transformer_class", "io.github.kavahub.learnjava.TransformerWithASM");
        Process process = builder.start();

        // 等待程序运行完成
        while(process.isAlive()) {}
    }

    @Test
    public void givenByteBuddy_whenRunProgream_thenCheckConsole() throws IOException, InterruptedException {
        // 启动程序
        ProcessBuilder builder = Main.build();
        builder.environment().put("transformer_class", "io.github.kavahub.learnjava.TransformerWithByteBuddy");
        Process process = builder.start();

        // 等待程序运行完成
        while(process.isAlive()) {}
    }
}

我们可以愉快的运行测试了，抛弃了繁琐的命令行。

全部的代码，可以在这里查看elapse-of-time ，欢迎顶赞，感谢！