Implementing AOP Method Timing with Java Instrumentation and ASM

Author: Zhang Chuchen, Java development engineer at Qunar since 2016, experienced with Android, RN, Angular and now focusing on data algorithm engineering.

Overview

This series consists of two articles divided into three parts: (1) introduction of Instrument‑related terminology and concepts; (2) Instrument code practice, replacing Spring AOP with a custom AOP implementation; (3) application of Instrument in Qunar's full‑link tracing system QTrace on the client side. By reading this series you will understand Instrument technology and how to apply it.

What is Instrument?

Instrument agent, also known as JPLISAgent (Java Programming Language Instrumentation Services Agent), is a service that provides support for Java bytecode instrumentation. It relies on the Attach API and JVMTI mechanisms.

What is Attach API?

Attach API is a tool for inter‑process communication between JVMs. It can send commands such as jstack, jps, jmap, or load a Java agent into another JVM process. For example, you can attach to a JVM with PID 1234 and load an agent JAR.

// VirtualMachine and related classes are in tools.jar
VirtualMachine vm = VirtualMachine.attach("1234");
try {
    vm.loadAgent(".../agent.jar"); // specify the agent JAR path
} finally {
    vm.detach();
}

What is JVMTI?

JVMTI (JVM Tool Interface) is a set of JVM‑exposed interfaces for extensions. It is event‑driven and allows developers to access thread information, class metadata, heap inspection, and various VM events, as well as set breakpoints and single‑step execution.

What is ASM?

ASM is a Java bytecode manipulation framework. It can generate or modify class files directly in binary form, either at load time or at runtime.

Proxy Implementations

cglib – a high‑performance code generation library that extends Java classes and implements interfaces at runtime; used by Hibernate for lazy loading and by Spring for dynamic proxies.

JDK dynamic proxy – requires an interface to be present.

Instrumentation Scenarios

Examples include BTrace (a safe dynamic tracing tool) and class hot‑deployment.

Stage 1: Using ASM to Modify Bytecode

The goal is to insert code that prints method execution time.

/**
 * Cost annotation definition
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface Cost {}

/**
 * Test class annotated with @Cost
 */
@Cost
public class TestCostTime {
    public void doTestFunction1() {
        System.out.println("do test function....");
    }
    @Cost
    public void doTestFunction2() {
        try {
            System.out.println("do test function 2....");
            Thread.sleep(1100);
        } catch (InterruptedException ex) {}
    }
}

AOP Invoker Definition

public abstract class AopInvoker {
    protected final String mClassName;
    protected final String mMethodName;
    public AopInvoker(String className, String methodName) {
        this.mClassName = className;
        this.mMethodName = methodName;
    }
    public abstract void aspectBeforeInvoke();
    public abstract void aspectAfterInvoke();
    public static AopInvoker newInvoker(String className, String methodName) {
        return new MethodCostAopInvoker(className, methodName);
    }
}

public class MethodCostAopInvoker extends AopInvoker {
    private Long mStartTime;
    public MethodCostAopInvoker(String className, String methodName) {
        super(className, methodName);
    }
    @Override
    public void aspectBeforeInvoke() {
        mStartTime = System.currentTimeMillis();
    }
    @Override
    public void aspectAfterInvoke() {
        long endTime = System.currentTimeMillis();
        System.out.println(String.format("%s.%s consume:%dms", mClassName, mMethodName, endTime - mStartTime));
    }
}

ASM Class Visitor

public class AopClassVisitor extends ClassVisitor {
    private final String mTargetPackageName;
    private String mActualClassFullName;
    private boolean mNeedModifyMethod = false;
    public AopClassVisitor(String targetPackageName, int api, ClassVisitor cv) {
        super(api, cv);
        this.mTargetPackageName = targetPackageName;
    }
    @Override
    public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
        super.visit(version, access, name, signature, superName, interfaces);
        mActualClassFullName = name.replace('/', '.');
        if (mActualClassFullName != null) {
            mNeedModifyMethod = mActualClassFullName.startsWith(mTargetPackageName);
        }
    }
    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
        if (mNeedModifyMethod) {
            return new AopMethodVisitor(api, mv, access, name, desc, mActualClassFullName, name);
        }
        return mv;
    }
    @Override
    public void visitEnd() {
        super.visitEnd();
        mNeedModifyMethod = false;
    }
}

ASM Method Visitor

public class AopMethodVisitor extends AdviceAdapter {
    private int mInvokerVarIndex = 0;
    private final String mClassName;
    private final String mMethodName;
    public AopMethodVisitor(int api, MethodVisitor originMV, int access, String methodName, String desc, String className, String methodName) {
        super(api, originMV, access, methodName, desc);
        this.mClassName = className;
        this.mMethodName = methodName;
    }
    @Override
    protected void onMethodEnter() {
        super.onMethodEnter();
        beginAspect();
    }
    @Override
    protected void onMethodExit(int opcode) {
        super.onMethodExit(opcode);
        afterAspect();
    }
    @Override
    public void visitMaxs(int maxStack, int maxLocals) {
        super.visitMaxs(maxStack + 2, maxLocals + 1);
    }
    private void beginAspect() {
        if (mv == null) return;
        mv.visitLdcInsn(mClassName);
        mv.visitLdcInsn(mMethodName);
        mv.visitMethodInsn(INVOKESTATIC, "asm/framework/AopInvoker", "newInvoker", "(Ljava/lang/String;Ljava/lang/String;)Lasm/framework/AopInvoker;", false);
        mv.visitInsn(DUP);
        mv.visitMethodInsn(INVOKEVIRTUAL, "asm/framework/AopInvoker", "aspectBeforeInvoke", "()V", false);
        mInvokerVarIndex = newLocal(Type.getType("Lasm/framework/AopInvoker;"));
        mv.visitVarInsn(ASTORE, mInvokerVarIndex);
    }
    private void afterAspect() {
        if (mv == null) return;
        mv.visitVarInsn(ALOAD, mInvokerVarIndex);
        mv.visitMethodInsn(INVOKEVIRTUAL, "asm/framework/AopInvoker", "aspectAfterInvoke", "()V", false);
    }
}

AOP Engine – Entry Point

public class AopEngine {
    public static byte[] doAspect(String targetPackageName, Class clazz) throws Exception {
        return doAspect(targetPackageName, new ClassReader(clazz.getName()));
    }
    public static byte[] doAspect(String targetPackageName, InputStream classInputStream) throws Exception {
        return doAspect(targetPackageName, new ClassReader(classInputStream));
    }
    public static byte[] doAspect(String targetPackageName, byte[] classBytes) throws Exception {
        return doAspect(targetPackageName, new ClassReader(classBytes));
    }
    private static byte[] doAspect(String targetPackageName, ClassReader classReader) throws Exception {
        ClassWriter classWriter = new ClassWriter(COMPUTE_MAXS);
        AopClassVisitor classVisitor = new AopClassVisitor(targetPackageName, ASM5, classWriter);
        classReader.accept(classVisitor, EXPAND_FRAMES);
        return classWriter.toByteArray();
    }
}

Test Code – ClassFile Bean

public class ClassFile {
    private Class clazz;
    private String classFilePath;
    private String packageName;
    public ClassFile(Class clazz, String classFilePath, String packageName) {
        this.clazz = clazz;
        this.classFilePath = classFilePath;
        this.packageName = packageName;
    }
    public String getPackageName() { return packageName; }
    public void setPackageName(String packageName) { this.packageName = packageName; }
    public Class getClazz() { return clazz; }
    public void setClazz(Class clazz) { this.clazz = clazz; }
    public String getClassFilePath() { return classFilePath; }
    public void setClassFilePath(String classFilePath) { this.classFilePath = classFilePath; }
    @Override
    public String toString() {
        return "ClassFile{" + "clazz=" + clazz + ", classFilePath='" + classFilePath + '\'' + ", packageName='" + packageName + '\'' + '}';
    }
}

Bytecode Modification Process

// Get classes annotated with @Cost
List
classList = PackageUtil.getAnnotationClasssFromPackage("rootPackage", Cost.class);
for (ClassFile className : classList) {
    File classFile = new File(className.getClassFilePath());
    System.out.println(className.getClazz().getName());
    try {
        // Modify bytecode
        byte[] tBytes = AopEngine.doAspect(className.getPackageName(), classFile);
        // Write back to file
        FileOutputStream fout = new FileOutputStream(classFile);
        fout.write(tBytes);
        fout.close();
    } catch (Exception e) {
        e.printStackTrace();
    }
}

After recompiling the modified class and running the test code, the console shows both the original output and the method execution time logged by the injected AOP logic.

(new TestCostTime()).doTestFunction2();
ClassLoader classLoader = TestCostTime.class.getClassLoader();
((TestCostTime)classLoader.loadClass("asm.demo.TestCostTime").newInstance()).doTestFunction2();

Conclusion

The current implementation modifies class files on disk; to achieve true runtime AOP you need to load the agent at JVM startup using Instrument technology, which will be covered in the second article of this series.