基于Java Agent的attach方式实现方法耗时监控

在上一篇中我们已经介绍了java agent的相关概念和思想，给出了premain方式的实现代码。本篇主要是实现了attach方式，不同之处主要如下:

1. premain是静态修改，在类加载之前修改; attach是动态修改，在类加载后修改
2. 要使premain生效重启应用，而attach不重启应用即可修改字节码并让其重新加载

可以看到attach的方式更加强大，其核心原理首先是找到相关的进程id, 然后根据进程id去动态修改相关字节码，具体的修改方式和premain无差，下面就直接给出详细实现。

项目结构(此处为了方便把主程序和Agent程序放在一起, 实际生产中肯定是分开的):

agentdemo
├── pom.xml
└── src
    └── main
        ├── java
        │   └── com
        │       └── hebh
        │           └── demo
        │               ├── agent
        │               │   ├── MyInstrumentationAgent.java
        │               │   └── MyTransformer.java
        │               └── application
        │                   ├── AgentLoader.java
        │                   ├── Launcher.java
        │                   ├── MyApplication.java
        │                   └── Runner.java
        └── resources
            ├── META-INF
            │   └── MANIFEST.MF
            └── log4j2.xml

先看测试结果:

1. 打成jar包

   mvn clean package

2. 运行主程序

   java -jar target/myAgent-jar-with-dependencies.jar

3. 运行agent程序, 注意带上系统的lib目录

   java -Djava.ext.dirs=${JAVA_HOME}/lib -jar target/myAgent-jar-with-dependencies.jar LoadAgent

   如下图所示，可以看到首先找到主程序的进程id为4477，然后再attach上去

image-20190126202600306

然后此时再看主程序的运行日志, 可以看到在attach后动态增加的字节码生效了，实现了方法耗时监控:

image-20190126203002824

详细代码:

MANIFEST.MF

Main-Class: com.hebh.demo.application.Launcher
Agent-Class: com.hebh.demo.agent.MyInstrumentationAgent
Can-Redefine-Classes: true
Can-Retransform-Classes: true

pom

    4.0.0
    com.hebh
    agent-demo
    jar
    1.0-SNAPSHOT
    A custom project using myfaces
    http://www.myorganization.org

    
        myAgent
        
            
                org.apache.maven.plugins
                maven-compiler-plugin
                
                    1.8
                    1.8
                
            
            
                org.apache.maven.plugins
                maven-assembly-plugin
                
                    
                        
                        src/main/resources/META-INF/MANIFEST.MF
                    
                    
                        
                        jar-with-dependencies
                    
                
                
                    
                        make-assembly 
                        package 
                        
                            single
                        
                    
                
            
        
    

    
    
        
            com.sun
            tools
            1.8
            system
            ${java.home}/../lib/tools.jar
        

        
        
            org.javassist
            javassist
            3.24.1-GA
        

        
            org.apache.logging.log4j
            log4j-api
            2.11.1
        
        
            org.apache.logging.log4j
            log4j-core
            2.11.1
        
    

主程序和Agent程序的路由:

public class Launcher {
    public static void main(String[] args) throws Exception {
        if(args != null && args.length > 0 && "LoadAgent".equals(args[0])) {
            new AgentLoader().run();
        }else{
            new MyApplication().run();
        }
    }
}

主程序部分:

public class MyApplication {
    private static Logger logger = LogManager.getLogger(MyApplication.class);

    public static void run() throws Exception {
        logger.info("[Application] Starting My application");
        Runner runner = new Runner();
        for(;;){
            runner.run();
        }
    }
}

public class Runner {
    private static final Logger logger = LogManager.getLogger(Runner.class);
    public void run() throws InterruptedException{
        long sleep = (long)(Math.random() * 1000 + 200);
        Thread.sleep(sleep);
        logger.info("run in [{}] millis!", sleep);
    }
}

Agent部分:

public class AgentLoader {
    private static Logger logger = LogManager.getLogger(AgentLoader.class);

    public static void run() {
        //指定jar路径
        String agentFilePath = "/Users/baohuahe/demos/agentdemo/target/myAgent-jar-with-dependencies.jar";
        
        //需要attach的进程标识
        String applicationName = "myAgent";

        //查到需要监控的进程
        Optional jvmProcessOpt = Optional.ofNullable(VirtualMachine.list()
                .stream()
                .filter(jvm -> {
                    logger.info("jvm:{}", jvm.displayName());
                    return jvm.displayName().contains(applicationName);
                })
                .findFirst().get().id());

        if(!jvmProcessOpt.isPresent()) {
            logger.error("Target Application not found");
            return;
        }
        File agentFile = new File(agentFilePath);
        try {
            String jvmPid = jvmProcessOpt.get();
            logger.info("Attaching to target JVM with PID: " + jvmPid);
            VirtualMachine jvm = VirtualMachine.attach(jvmPid);
            jvm.loadAgent(agentFile.getAbsolutePath());
            jvm.detach();
            logger.info("Attached to target JVM and loaded Java agent successfully");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}

public class MyInstrumentationAgent {
    private static Logger logger = LogManager.getLogger(MyInstrumentationAgent.class);

    public static void agentmain(String agentArgs, Instrumentation inst) {
        logger.info("[Agent] In agentmain method");

        //需要监控的类
        String className = "com.hebh.demo.application.Runner";
        transformClass(className, inst);
    }

    private static void transformClass(String className, Instrumentation instrumentation) {
        Class targetCls = null;
        ClassLoader targetClassLoader = null;
        // see if we can get the class using forName
        try {
            targetCls = Class.forName(className);
            targetClassLoader = targetCls.getClassLoader();
            transform(targetCls, targetClassLoader, instrumentation);
            return;
        } catch (Exception ex) {
            logger.error("Class [{}] not found with Class.forName");
        }
        // otherwise iterate all loaded classes and find what we want
        for(Class clazz: instrumentation.getAllLoadedClasses()) {
            if(clazz.getName().equals(className)) {
                targetCls = clazz;
                targetClassLoader = targetCls.getClassLoader();
                transform(targetCls, targetClassLoader, instrumentation);
                return;
            }
        }
        throw new RuntimeException("Failed to find class [" + className + "]");
    }

    private static void transform(Class clazz, ClassLoader classLoader, Instrumentation instrumentation) {
        MyTransformer dt = new MyTransformer(clazz.getName(), classLoader);
        instrumentation.addTransformer(dt, true);
        try {
            instrumentation.retransformClasses(clazz);
        } catch (Exception ex) {
            throw new RuntimeException("Transform failed for class: [" + clazz.getName() + "]", ex);
        }
    }

}

public class MyTransformer implements ClassFileTransformer {
    private static Logger logger = LogManager.getLogger(MyTransformer.class);

    //需要监控的方法
    private static final String WITHDRAW_MONEY_METHOD = "run";

    /** The internal form class name of the class to transform */
    private String targetClassName;
    /** The class loader of the class we want to transform */
    private ClassLoader targetClassLoader;

    public MyTransformer(String targetClassName, ClassLoader targetClassLoader) {
        this.targetClassName = targetClassName;
        this.targetClassLoader = targetClassLoader;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] byteCode = classfileBuffer;
        String finalTargetClassName = this.targetClassName.replaceAll("\\.", "/"); //replace . with /
        if (!className.equals(finalTargetClassName)) {
            return byteCode;
        }

        if (className.equals(finalTargetClassName) && loader.equals(targetClassLoader)) {
            logger.info("[Agent] Transforming class" + className);
            try {
                ClassPool cp = ClassPool.getDefault();
                CtClass cc = cp.get(targetClassName);
                CtMethod m = cc.getDeclaredMethod(WITHDRAW_MONEY_METHOD);

                // 开始时间
                m.addLocalVariable("startTime", CtClass.longType);
                m.insertBefore("startTime = System.currentTimeMillis();");

                StringBuilder endBlock = new StringBuilder();

                // 结束时间
                m.addLocalVariable("endTime", CtClass.longType);
                endBlock.append("endTime = System.currentTimeMillis();");

                // 时间差
                m.addLocalVariable("opTime", CtClass.longType);
                endBlock.append("opTime = endTime-startTime;");

                // 打印方法耗时
                endBlock.append("logger.info(\"completed in:\" + opTime + \" millis!\");");

                m.insertAfter(endBlock.toString());

                byteCode = cc.toBytecode();
                cc.detach();
            } catch (Exception e) {
                logger.error("Exception", e);
            }
        }
        return byteCode;
    }
}