字节码编程ASM之插桩方法调用记录

写在前面

源码。
正式开始之前，先分享一个让人”悲伤“的真实的故事。
那是一个风和日丽的周六的下午，俺正在开开心心的打着羽毛球，突然接到了来自于最不想联系的那个人（没错，这个人就是我的领导！！！）的微信语音，语音里领导说线上出现问题：工作流任务办理，点击办理，不出办理人了，你到明天去公司加班看下咋回事吧！。瞬间整个人都不好了啊。

事后复盘导致bug出现的原因是测试场景覆盖的不够全导致（不要脸，代码写的有问题，就知道推给可爱的测试同学），也就是导致出现问题的方法在测试环境压根就没有调用到过。

而且因为这种原因导致的线上bug，已然不是第一次遇到了，所以我就在想，有没有什么办法能够规避这种问题的发生呢？测试用例每次都会投入大量的时间和测试同学一起评审啊！还是会出现场景覆盖不全的问题。头痛医头，脚痛医脚，既然是方法没有调用到，那我们只要在方法调用的时候打个日志记录下，测试后看下哪个方法没有日志不就是没有调用到吗！但额外的去写这种侵入业务的代码感觉不是太好，最后，想到了字节码插桩技术，只要在方法执行前插入一行日志打印代码，再和javaagent结合在一起使用不就行了，对程序完全无侵入。

嗯，完美的方案！！！

1：ASM版本

我们一步步来，首先要搞清楚程序的执行顺序，然后再搞清楚核心步骤都做了哪些事情，最后最好再自己手撕一遍。
首先我们要来定义premain，这是javaagent执行的入口：

/**
 * premain 程序，清单文件中需要配置的入口类
 */
public class PreMain {
    //JVM 首先尝试在代理类上调用以下方法
    public static void premain(String agentArgs, Instrumentation inst) {
        inst.addTransformer(new MethodRecordingTransformer());
    }

    //如果代理类没有实现上面的方法，那么 JVM 将尝试调用该方法
    public static void premain(String agentArgs) {
    }
}

类MethodRecordingTransformer是自定义的转换器，通过addTransformer方法注册到Instrument对象，这个转换器实现了转换规范接口java.lang.instrument.ClassFileTransformer,用来将字节码从一种形式转换为另外一种形式（这个过程就是插桩了）。如下：

public class MethodRecordingTransformer implements ClassFileTransformer {

    /**
     * 所有的类被来加载器加载的时候都会走这个方法，这样我们就有机会通过各种技术手段来对字节码进行插桩了
     *
     * @param loader
     * @param className           类全限定名称
     * @param classBeingRedefined
     * @param protectionDomain
     * @param classfileBuffer     类字节码对应的二进制数组，可以通过defineClass直接加载到JVM中并生成Class对象
     * @return
     * @throws IllegalClassFormatException
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            // 过滤不需要插桩的类
            if (MethodRecordingFilter.isNotNeedInject(className)) {
                return classfileBuffer;
            }
            // 插桩
            byte[] bytes = generateInstrumentationCode(loader, className, classfileBuffer);
            // 写到磁盘的.class文件中，方便调试
            outputClazz(bytes, className);
            return bytes;
        } catch (Throwable e) {
            System.out.println(e.getMessage());
        }
        return classfileBuffer;

    }

    private void outputClazz(byte[] bytes, String className) {
        // 输出类字节码
        FileOutputStream out = null;
        try {
            String pathName = MethodRecordingTransformer.class.getResource("/").getPath() + className + "SQM.class";
            out = new FileOutputStream(new File(pathName));
            System.out.println("插桩后代码输出路径：" + pathName);
            out.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != out) try {
                out.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 生成插桩代码
     *
     * @param loader
     * @param className
     * @param classfileBuffer
     * @return
     */
    private byte[] generateInstrumentationCode(ClassLoader loader, String className, byte[] classfileBuffer) {
        // 读取原有类
        ClassReader cr = new ClassReader(classfileBuffer);
        // 通过writer向需要插桩类写入插桩逻辑
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
        // 使用子自定义的类访问器来定义插桩逻辑
        ClassVisitor cv = new MethodRecordingClassAdapter(cw, className);
        // 真正插桩
        cr.accept(cv, ClassReader.EXPAND_FRAMES);
        return cw.toByteArray();
    }
}

主要看法方法generateInstrumentationCode，用来生成插桩后的二进制字节码文件，也就是真正被加载到JVM中运行的代码。generateInstrumentationCode方法中使用MethodRecordingClassAdapter，这是自定义的ClassVisitor的子类，用来进行判断是否需要过滤等操作，如下：

public class MethodRecordingClassAdapter extends ClassVisitor {
    // snip code

    @Override
    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
        // 不对接口和私有方法注入
        if (isInterface || (access & ACC_PRIVATE) != 0) {
            return super.visitMethod(access, name, descriptor, signature, exceptions);
        }
        //不对抽象方法、native方法、桥接方法、合成方法进行注入
        if ((access & ACC_ABSTRACT) != 0
                || (access & ACC_NATIVE) != 0
                || (access & ACC_BRIDGE) != 0
                || (access & ACC_SYNTHETIC) != 0) {
            return super.visitMethod(access, name, descriptor, signature, exceptions);
        }
        if ("<init>".equals(name) || "<clinit>".equals(name)) {
            return super.visitMethod(access, name, descriptor, signature, exceptions);
        }
        // 过滤Object类默认方法
        if (MethodRecordingFilter.isNotNeedInjectMethod(name)) {
            return super.visitMethod(access, name, descriptor, signature, exceptions);
        }
        // 获取要增强方法的MethodVisitor对象
        MethodVisitor mv = cv.visitMethod(access, name, descriptor, signature, exceptions);
        if (null == mv) return null;
        // 使用扩展的MethodVisitor来对方法进行插桩
        return new MethodRecordingVisitorAdvice(access, name, descriptor, mv, className, fullClazzName, simpleClassName);
    }
}

主要看方法visitMethod最终返回的MethodVisitor类型的对象MethodRecordingVisitorAdvice，定义了方法执行前以及方法返回前要执行的逻辑：

public class MethodRecordingVisitorAdvice extends AdviceAdapter {
    private final String className;
    private String methodName;

    protected MethodRecordingVisitorAdvice(int access, String methodName, String desc, MethodVisitor mv, String className, String fullClassName, String simpleClassName) {
        super(ASM5, mv, access, methodName, desc);
        this.methodName = methodName;
        this.className = className;
    }

    @Override
    protected void onMethodEnter() {
        // public static void printLog(String className, String methodName) {
        mv.visitLdcInsn(className);
        mv.visitLdcInsn(methodName);
        mv.visitMethodInsn(INVOKESTATIC, "com/dahuyou/method/record/LogUtil", "printLog", "(Ljava/lang/String;Ljava/lang/String;)V", false);
    }

    /**
     * 方法执行结束前调用，即return，throw异常前调用
     *
     * @param opcode 操作码 org.objectweb.asm.Opcodes 比如方法的public private啊，jdk的版本号啊，位运算啊，各种字节码指令啊当然包含这个方法会传入的如下可能的操作码：
     * int IRETURN = 172; // 返回int
     * int LRETURN = 173; // 返回long
     * int FRETURN = 174; // 返回float
     * int DRETURN = 175; // 返回double
     * int ARETURN = 176; // 返回引用类型
     * int RETURN = 177; // 返回void
     * 这里就可以通过opcode知道方法返回的值是什么了！
     */
    @Override
    protected void onMethodExit(int opcode) {
        super.onMethodExit(opcode);
    }
}

这里只对onMethodEnter方法添加了插桩代码，即方法执行前调用LogUtil的静态方法printLog:

public class LogUtil {
    public static void printLog(String className, String methodName) {
        // 实际场景根据需要就可以写到日志文件，或者是对接到消息中间件了
        System.out.println("类：[" + className + "], 方法：[" + methodName + "] 被测试同学覆盖到了！");
    }
}

这里只是模拟了，实际场景根据具体情况修改LogUtil类就行了。

然后我们就可以来测试了，首先来打包：
在这里插入图片描述
准备测试代码：

/*
获取工作流任务办理人的工具类
 */
public class WorkflowUtil {
    public List<String> queryAssignee() {
        return null;
    }
}

public class Tst {
    public static void main(String[] args) {

        new WorkflowUtil().queryAssignee();
    }
}