如何编写JavaAgent

这篇文章是根据MegaEase的袁伟老师的分享而来,地址是How To Write a JavaAgent

简介

java agent是什么？

java agent是jdk1.5时候推出的一个在运行时动态修改class,从而达到动态修改行为的目的

能做什么？

功能与AOP类似，它的优势在与彻底和业务代码隔离，可以完成AOP相同的事情，并且不入侵业务代码，适合于日志采集、链路追踪等基础组件

用法

java agent主要可以在两个时间点进行加载：

1. JVM启动时
2. 目标方法运行时

项目结构

启动时加载示例

• AgentExampleDemo

public class AgentExampleDemo {
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        System.out.println("premain start...");
    }
}

• AgentTarget

public class AgentTarget {
  public static void main(String[] args) {
      System.out.println("开始打印....");
  }
}

• build.gradle

  //重新定义MANIFEST.MF
  jar {
    manifest {
        attributes 'Premain-Class': 'com.agmtopy.source.agent.AgentExampleDemo'
    }
}

• javaagent启动参数

-javaagent:build/libs/jvmsource-1.0-SNAPSHOT.jar

• 运行结果

这是第一种使用agent的方式,在目标代码运行前使用,java agent代码与目标方法进行组合的方式进行执行

运行时加载示例

目标方法远行时加载要使用到javassist这个工具帮助我们修改class,注意javassist有两个项目,要使用org.javassist😂才可以

implementation "org.javassist:javassist:3.28.0-GA"

1. 修改AgentTarget

保持jvm运行,以便通过Attach的方式进行替换

public class AgentTarget {
    public static void main(String[] args) throws Exception {
        System.out.println("开始打印....");
        while (true) {
            TimeUnit.SECONDS.sleep(5);
            println();
        }
    }

    public static void println() {
        System.out.println(LocalTime.now());
    }
}

2. 在AgentExampleDemo中增加agentmain方法

• AgentExampleDemo

/**
 * attach:方式运行
 */
public static void agentmain(String agentArgs, Instrumentation inst) {
    System.out.println("agentmain start...");
    // 显示执行时间
    inst.addTransformer(new ShowExecTime(), true);
    try {
        //重写载入新的字节码
        inst.retransformClasses(AgentTarget.class);
    } catch (UnmodifiableClassException e) {
        e.printStackTrace();
    }
}

3. 增加ShowExecTime来修改字节码

• ShowExecTime

/**
* 自定义ClassFileTransformer
*/
public class ShowExecTime implements ClassFileTransformer {

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        //只针对目标包下进行耗时统计
        if (!className.startsWith("com/agmtopy/source/agent")) {
            return classfileBuffer;
        }
        System.out.println("正在加载类：" + className);

        try {
            ClassPool classPool = ClassPool.getDefault();
            classPool.appendClassPath(new LoaderClassPath(loader));

            CtClass cl = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
            // 所有方法，统计耗时
            for (CtMethod method : cl.getDeclaredMethods()) {
                System.out.println("开始修改:" + method +" 方法" );
                //需要通过`addLocalVariable`来声明局部变量
                method.addLocalVariable("start", CtClass.longType);
                //插入 开始语句
                method.insertBefore("start = java.lang.System.currentTimeMillis();");
                String methodName = method.getLongName();
                //创建并插入 打印语句 System.out.println("方法：test， 执行时间：" + (System.currentTimeMillis() - start));
                String statement = String.format("java.lang.System.out.println(\"方法：%s， 执行时间：\" + (java.lang.System.currentTimeMillis() - start));", methodName);
                method.insertAfter(statement);
            }

            byte[] transformed = cl.toBytecode();
            return transformed;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }
}

transform()方法是通过javassist来修改字节码,在方法执行前后插入局部变量，然后打印方法执行耗时

4. 运行jar进行替换字节码替换

/**
 * 修改指定运行中的代码
 */
public static void main(String[] args) throws Exception {
    // 传入目标 JVM pid
    VirtualMachine vm = VirtualMachine.attach("6068");
    vm.loadAgent("D:\\project\\jvmsource\\build\\libs\\jvmsource-1.0-SNAPSHOT.jar");
}

5. 修改MANIFEST.MF
   需要将Agent-Class写入MANIFEST.MF文件

• gradle

jar {
    manifest {
        attributes 'Can-Redefine-Classes': true
        attributes 'Can-Retransform-Classes': true
        attributes 'Agent-Class': 'com.agmtopy.source.agent.AgentExampleDemo'
        attributes 'Premain-Class': 'com.agmtopy.source.agent.AgentExampleDemo'
    }
}

6. 执行字节码替换

6.1 先将项目构建成为jar包
6.2 运行AgentTarget
不需要使用-javaagent的方式进行启动

6.3 执行字节码替换

• 运行结果
  

源码分析

java agent的原理根据加载时机还是可以分为两类入口，一类是启动时将agent class挂载到目标JVM上，另外一类入口是运行时动态加载，采用的是JVM attach技术

启动时加载原理分析

分析目标方法调用链

public class AgentExampleDemo {

    public static void premain(String agentArgs, Instrumentation instrumentation) {
        System.out.println("premain start...");
        println(Thread.currentThread().getStackTrace());
    }

    //打印调用栈
    public static void println(StackTraceElement[] elements) {
        for (int i = 0; i < elements.length; i++) {
            StringBuffer buffer = new StringBuffer();
            buffer.append("index: ").append(i).append(" ClassName: ").append(elements[i].getClassName())
                    .append(" Method Name : " + elements[i].getMethodName());
            System.out.println(buffer.toString());
        }
    }
}

通过调用栈可以分析出是InstrumentationImpl调用premain()方法的,下面开始分析InstrumentationImpl

InstrumentationImpl

public class InstrumentationImpl implements Instrumentation 

InstrumentationImpl实现Instrumentation,Instrumentation接口是JVM定义对字节码操作的接口，我们按照调用链的顺序倒叙进行分析（执行、触发）

1. permain执行过程分析

由于InstrumentationImpl.loadClassAndCallPremain()方法已经最顶层的java代码入口，通过方法名称查找可以在JPLISAgent.h文件中查询到该方法名称被定义成为一个常量

• JPLISAgent.h

#define JPLIS_INSTRUMENTIMPL_PREMAININVOKER_METHODNAME      "loadClassAndCallPremain"

该常量被JPLISAgent.c的createInstrumentationImpl方法所使用

• JPLISAgent.c

jboolean createInstrumentationImpl(JNIEnv *jnienv,JPLISAgent *agent)
{
    //省略....
    /* Now look up the method ID for the pre-main caller (we will need this more than once) */
    if (!errorOutstanding)
    {
        //①获取到调用permain方法MethodId
        premainCallerMethodID = (*jnienv)->GetMethodID(jnienv,
                                                       implClass,
                                                       JPLIS_INSTRUMENTIMPL_PREMAININVOKER_METHODNAME,
                                                       JPLIS_INSTRUMENTIMPL_PREMAININVOKER_METHODSIGNATURE);
    }

    if (!errorOutstanding)
    {
        agent->mInstrumentationImpl = resultImpl;
        //②指针赋值
        agent->mPremainCaller = premainCallerMethodID;
        agent->mAgentmainCaller = agentmainCallerMethodID;
        agent->mTransform = transformMethodID;
    }
    
    //省略....
    return !errorOutstanding;
}

这段方法主要是做两件事，第一是获取到调用permain方法MethodId；第二件事是将这个MethodId传递出去,分析mPremainCaller的使用可以得到该值在’processJavaStart’中使用

• processJavaStart

result = startJavaAgent(agent, jnienv,
                        agent->mAgentClassName, 
                        agent->mOptionsString,
                        agent->mPremainCaller);

processJavaStart通过前面获取到的MethodId启动javaAgent，下面我们分析statrJavaAgent

• statrJavaAgent

//...
success = invokeJavaAgentMainMethod(jnienv,
                                    agent->mInstrumentationImpl,
                                    agentMainMethod,
                                    classNameObject,
                                    optionsStringObject);

调用invokeJavaAgentMainMethod传入对象、方法ID、实参，执行定义的premain方法

通过上面的分析知道了permai执行的过程，继续看一下permain方法是如何触发的

2. permain触发过程分析

processJavaStart方法是执行permain的入口，它在JPLISAgent.h中进行定义的,在源代码中全局搜索： JPLISAgent *可以找到JPLISAgent是在<InvocationAdapter.c>中重新进行过赋值

• InvocationAdapter.c

void JNICALL eventHandlerVMInit(jvmtiEnv *jvmtienv,
                    JNIEnv *jnienv,
                    jthread thread)
  {
      JPLISAgent *agent = environment->mAgent;
  }

eventHandlerVMInit方法在JPLISAgent.c的initializeJPLISAgent方法中被设置为回调方法

• initializeJPLISAgent

//关键执行逻辑
if (jvmtierror == JVMTI_ERROR_NONE)
{
    jvmtiEventCallbacks callbacks;
    memset(&callbacks, 0, sizeof(callbacks));
    //设置JVM回调
    callbacks.VMInit = &eventHandlerVMInit;

    jvmtierror = (*jvmtienv)->SetEventCallbacks(jvmtienv,
                                                &callbacks,
                                                sizeof(callbacks));
    check_phase_ret_blob(jvmtierror, JPLIS_INIT_ERROR_FAILURE);
    jplis_assert(jvmtierror == JVMTI_ERROR_NONE);
}

通过这里可以看到在初始化JPLISAgent时候就设置了JVM初始化完成后会回调InvocationAdapter来执行permain方法

3. JPLISAgent的初始化

回溯initializeJPLISAgent方法可以找到分别在InvocationAdapter.c的DEF_Agent_OnLoad、DEF_Agent_OnAttach上被调用。这两种方式也正是前面讲到的agent的两种增强方式的入口。

在JVM启动时最开始加载的是libinstrument动态链接库，然后在动态链接库里面找到JVMTI的入口方法：Agent_OnLoad和Agent_OnAttach。InvocationAdapter.c的定义

• InvocationAdapter.c

/*
*  This will be called once for every -javaagent on the command line.
*  Each call to Agent_OnLoad will create its own agent and agent data.
*
*  The argument tail string provided to Agent_OnLoad will be of form
*  <jarfile>[=<options>]. The tail string is split into the jarfile and
*  options components. The jarfile manifest is parsed and the value of the
*  Premain-Class attribute will become the agent's premain class. The jar
*  file is then added to the system class path, and if the Boot-Class-Path
*  attribute is present then all relative URLs in the value are processed
*  to create boot class path segments to append to the boot class path.
*/
JNIEXPORT jint JNICALL
DEF_Agent_OnLoad(JavaVM *vm, char *tail, void *reserved)

/*
*  This will be called once each time a tool attaches to the VM and loads
*  the JPLIS library.
*/
JNIEXPORT jint JNICALL
DEF_Agent_OnAttach(JavaVM *vm, char *args, void *reserved)

4. 整体逻辑
   整体的执行逻辑就是：

• JPLISAgent声明JVM启动时候初始化JPLISAgent
• JPLISAgent初始化时设置InvocationAdapter的回调方法
• JVM初始化完成后执行回调方法
• InvocationAdapter的回调方法执行permain方法

运行时加载原理分析

1. 分析字节码
   使用HSDB查看AgentTarget未进行字节码替换前的数据

jhsdb hsdb --pid 21656

• 常量池

明显可以看到常量池中增加了22条指令，这22条指令就是新加入的字节码所要使用到的常量

• 方法区

2. 重新加载字节码原理
   @TODO

通过字节码对比可以明显的看出在使用Instrumentation.addTransformer();后确实将字节码进行了修改,其实修改字节码还是有两个时机:

• 一个是在启动编译时,会在jvm启动完毕后在执行permain方法来修改字节码
• 一个就是在运行期间动态的修改字节码

整体执行逻辑

开发工具

• ASM

• Javassist

• Byte Buddy

功能对比

-	ASM	Javassist	Byte Buddy
学习成本	高	低	低
使用方法	使用字节码方式进行插入,需要了解class类结构和JVM指令集	提供高级抽象接口和低级字节码接口	同Javassist,并且提供声明式接口
性能	极快	一般	快

详细的对比可参考byteBuddy官方资料:https://bytebuddy.net/#/tutorial

Byte Buddy

示例

• gradle

//目前最新版本为1.11.6
implementation "net.bytebuddy:byte-buddy:LATEST"

• AgentExampleForByteBuddy

public class AgentExampleForByteBuddy {

    @Advice.OnMethodEnter
    public static boolean before(@Advice.Origin Executable method) {
        System.out.println("byte buddy before : " + method);
        return true;
    }

    @Advice.OnMethodExit
    public static void after(@Advice.Origin Executable method) {
        System.out.println("byte buddy after : " + method);
    }

    public static void premain(String arguments, Instrumentation inst) {
        System.out.println("开始执行...");
        new AgentBuilder.Default()
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.InstallationListener.StreamWriting.toSystemError())
                .type(ElementMatchers.nameContains("AgentTarget"))
                .transform((builder, td, cl, m) -> builder.visit(Advice.to(AgentExampleForByteBuddy.class).on(MethodDescription::isConstructor)))
                .installOn(inst);
    }
}

• AgentTarget

public class AgentTarget {
    public static void main(String[] args) throws Exception {
        System.out.println("开始打印....");
        AgentTarget agentTarget = new AgentTarget();
        agentTarget.printInfo();
    }

    public void printInfo() {
        System.out.println("123" + LocalTime.now());
    }
}

这里演示的是一个重写加载类的示例:

1. 通过@Advice.OnMethodEnter和@Advice.OnMethodExit定义执行方法前后插入的字节码
2. 通过AgentBuilder指定要增强的类和类型
3. AgentBuilder.with可以添加监听，方便输出调试

Byte Byddy常见问题

• 依赖冲突

处理方案：

1. 构建工具排除

2. 使用自定义classLoader加载agent所使用的类
   通过指定agent类的加载器,让BootstrapClassLoader去加载

• Byte-Buddy提供的API

ClassInjector.UsingInstrumentation
         .of(FileJar, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, instrumentation)
         .injectRaw(Collections.singletonMap(instrumentation, ClassFileLocator
                 .ForClassLoader.ofSystemLoader()
                 .locate("com.agmtopy.source.agent.AgentExampleForByteBuddy").resolve()));

• java agent中传参

处理方案:

1. ThreadLocal
2. 增加临时的成员变量

参考资料

https://asm.ow2.io/
https://www.javassist.org/
https://bytebuddy.net/#/
https://blog.csdn.net/wanxiaoderen/article/details/107079741
https://www.cnblogs.com/old-cha/p/13264114.html
https://www.cnblogs.com/chiangchou/p/javassist.html#_label9
https://tech.meituan.com/2019/11/07/java-dynamic-debugging-technology.html
https://github.com/gzzchh/de-ag
https://gitee.com/mazhimazh/bytecode-examples