字节码框架实践

字节码框架创建Agent实践

目标

使用不同字节码框架构建JavaAgent，并实现方法运行耗时监控

测试类

App作为入口类，调用Foo的相关方法打印“Hello world!”

public class App {
    public static void main( String[] args ) throws InterruptedException {
        Foo foo = new Foo();
        foo.echo();
        foo.shutdown();
    }
}

/**
 * 测试类，两个方法，简单打印helloworld
 */
public class Foo {
    public void echo() {
        System.out.println("Hello world!");
        try {
            Thread.sleep((long) (Math.random() * 100));
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public void shutdown() {
        System.out.println("hello world again!");
    }
}

打包配置

不同的字节码框架实现方案中，打包agent的配置基本一样，先列出来。

由于Agent对字节码的修改依赖了相关Jar包，因此打包时需要使用assembly插件打一个fatjar，将所有依赖的jar包打进Agent jar中。否则会出现运行失败。

<plugin>
  <groupId>org.apache.maven.plugins</groupId>
  <artifactId>maven-assembly-plugin</artifactId>
  <executions>
    <execution>
      <goals>
        <goal>attached</goal>
      </goals>
      <phase>package</phase>
      <configuration>
        <descriptorRefs>
          <descriptorRef>jar-with-dependencies</descriptorRef>
        </descriptorRefs>
        <archive>
          <manifestEntries>
            <Premain-class>com.example.asm.agent.AsmAgent</Premain-class>
<!--            <Premain-class>com.example.javaassist.AssistAgent</Premain-class> -->
<!--            <Premain-class>com.example.bytebuddy.ByteBuddyAgent</Premain-class> -->
            <Can-Redefine-Classes>true</Can-Redefine-Classes>
            <Can-Retransform-Classes>true</Can-Retransform-Classes>
          </manifestEntries>
        </archive>
      </configuration>
    </execution>
  </executions>
</plugin>

ASM实现

Agent

实现premain方法绑定ClassFileTransformer ，后者负责在transform阶段使用ASM框架完成字节码操作。

ASM的API使用访问者模式，需要创建ClassReader ClassWriter ClassVisitor 分别负责读取Class字节码，写Class字节码，操作Class字节码。增强逻辑在ClassVisitor中实现。

public class AsmAgent {
    public static void premain(String premainArgs, Instrumentation instrumentation) {
        // 设置classfileTransformer
        instrumentation.addTransformer(new AsmClassFileTransformer());

    }

    // asm的实现，完成对class字节码的修改
    static class AsmClassFileTransformer implements ClassFileTransformer {

        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            try {
                if (className.contains("Foo")) {
                    // 读Class字节码
                    ClassReader classReader = new ClassReader(className);
                    ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_MAXS);
                    ClassVisitor classVisitor = new ApmClassVisitor(classWriter);
                    // 处理class字节码
                    classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);

                    // 打印到本地（主要用于观察学习）
                    byte[] arr = classWriter.toByteArray();
                    File f = new File("/Users/yuangong/IdeaProjects/agentlearn/src/main/java/com/example/asm/Foo1.class");
                    FileOutputStream fos = new FileOutputStream(f);
                    fos.write(arr);
                    fos.close();
                    return arr;
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
            return classfileBuffer;
        }
    }
}

Visitor增强逻辑

ApmClassVisitor,ApmMethodVistitor逻辑如下，二者在visitor访问到具体字节码结构时，执行增强逻辑

public class ApmClassVisitor extends ClassVisitor implements Opcodes {
    public ApmClassVisitor(ClassVisitor classVisitor) {
        super(ASM5, classVisitor);
    }

    /**
     * 不处理初始化、构造方法
     */
    @Override
    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
        MethodVisitor methodVisitor = cv.visitMethod(access, name, descriptor, signature, exceptions);
        if (!name.equals("<init>") && methodVisitor != null) {
            methodVisitor = new ApmMethodVisitor(methodVisitor, access, name, descriptor);
        }
        return methodVisitor;
    }
} 

public class ApmMethodVisitor extends AdviceAdapter {
    public ApmMethodVisitor(MethodVisitor mv, int access, String name, String desc) {
        super(ASM5, mv, access, name, desc);
    }

    // visitor访问到方法进入处回调，添加字节码
    @Override
    protected void onMethodEnter() {
        // 添加字节码：调用指定类的方法
        mv.visitMethodInsn(INVOKESTATIC, "com/example/asm/agent/Apm", "before", "()V", false);
    }

    // visitor访问到方法退出处回调，添加字节码
    @Override
    protected void onMethodExit(int opcode) {
        // 添加字节码：调用指定类的方法
        mv.visitMethodInsn(INVOKESTATIC, "com/example/asm/agent/Apm", "end", "()V", false);
    }
}

APM是为了简化字节码操作编码困难封装的一个类，封装了耗时计算逻辑

public class Apm {
    private static long startTime = 0L;
    public static void before() {
        startTime = System.nanoTime();
        System.out.println("startTime = " + startTime);
    }

    public static void end() {
        long endTime = System.nanoTime();
        System.out.println("endTime = " + endTime);
        long cost = (endTime - startTime) / 1000000L;
        System.out.println("cost = " + cost);
    }
}

打包运行

按照打包配置，配置premain运行的类。打包生成jar包后，通过Idea执行启动参数运行即可。

ASM Bytecode Outline插件

由于字节码的学习成本很高，开发效率低。所以ASM开发了一个插件，可以直接在Idea中将Java代码转换成了对应的ASM代码，这样可以提高一些开发效率。安装时，在Idea的plugins中查询即可。使用方式如下：

创建一个原始类
右键，选择【show bytecode outline】，可以看到相关代码
增加代码，再次查看ASM代码，点击【show differences】，不同的代码即是需要添加到Visitor中的内容

Javaassist实现

Agent

实现premain方法，设置ClassFileTransformer

public class AssistAgent {
    public static void premain(String preArgs, Instrumentation instrumentation) {
        instrumentation.addTransformer(new AssistClassFileTransformer());
    }
}

类增强逻辑

AssistClassFileTransformer 中使用JavaAssist对字节码进行增强。JavaAssist主要使用CtClass,CtMethod两个类。分别表示获取的Class信息和Method信息。ClassPool 是类信息池，可以通过类名获取到CtClass 信息

public class AssistClassFileTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        if (className.contains("com/example/Foo")) {
            try {
                // 获取指定的类和方法
                CtClass ctClass = ClassPool.getDefault().get(className.replaceAll("/", "."));
                CtMethod ctMethod = ctClass.getDeclaredMethod("shutdown");
                // 额外：除了增加方法，还可以修改shutdown代码
                ctMethod.setBody( "System.out.println(\"javaassist goodbye!\");");

                for (CtMethod method : ctClass.getDeclaredMethods()) {
                    // 给所有方法添加监控
                    method.insertBefore("com.example.asm.agent.Apm.before();");
                    method.insertAfter("com.example.asm.agent.Apm.end();");
                }

                // 输出转换后的字节码，用于观察学习
                ctClass.writeFile("/Users/yuangong/IdeaProjects/agentlearn/src/main/java/com/example/javaassist");
                return ctClass.toBytecode();
            } catch (NotFoundException | CannotCompileException | IOException e) {
                e.printStackTrace();
            }
        }
        return classfileBuffer;
    }
}

打包运行

打包运行逻辑如ASM，不再赘述

ByteBuddy实现

Agent

ByteBuddy的Agent主要逻辑并没有直接设置Transformer，而是通过ByteBuddy特定的API生成并设置transformer。

public class ByteBuddyAgent {
    public static void premain(String premainArgs, Instrumentation instrumentation) {
        // 创建transformer
        AgentBuilder.Transformer transformer = new AgentBuilder.Transformer() {
            @Override
            public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule) {
                return builder.method(ElementMatchers.any()) // 哪些方法需要增强
                        .intercept(MethodDelegation.to(ApmComputer.class)); // 指定过滤增强的代理方法
            }
        };

        new AgentBuilder.Default()
                .type(ElementMatchers.nameStartsWith("com.example.Foo")) // 指定哪些类需要增强
                .transform(transformer).installOn(instrumentation); // 设置transformer
    }
}

增强逻辑

增强逻辑实现一个intercept方法，可以使用ByteBuddy的注解，指定获取响应的信息。

@RuntimeType 表示运行时定义的目标方法

@Origin 表示原始方法信息

@SuperCall 用于调用父类版本的方法

public class ApmComputer {
    @RuntimeType
    public static Object intercept(@Origin Method method, @SuperCall Callable<?> callable) throws Exception {
        long start = System.nanoTime();
        try {
            // 调用原有逻辑
            return callable.call();
        } finally {
            long end = System.nanoTime();
            long cost = (end - start) / 1000;
            System.out.println(method.getDeclaringClass().getName() + "." + method.getName() + " cost = " + cost);
        }
    }
}