javaagent 实现方法执行时间监控

小知识，大挑战！本文正在参与“程序员必备小知识”创作活动。

前言

闲来无事想想skywalking 是如何实现链路时长计算的呢，比如打印sql的执行时长的。skywaylking 的探针包 skywalking-agent.jar 是在项目启动时通过 -javaagent: 参数注入的，以此揭开学习agent 机制的序幕

在JDK1.5以后，我们可以使用agent技术构建一个独立于应用程序的代理程序（即为Agent），用来协助监测、运行甚至替换其他JVM上的程序。使用它可以实现虚拟机级别的AOP功能。

这么讲有些枯燥，用一个具体的demo来看看

思路

就拿打印sql执行时长来说，其实最简单的思路就是在方法执行前记录时间，执行后计算差值，比如

{

    long l = System.currentTimeMillis();
    // method body
    
    System.out.println("执行完毕用时：" + (System.currentTimeMillis() - l));

}

复制代码

以mybatis为例所有sql执行必须经过类org.apache.ibatis.executor.statement.PreparedStatementHandler,不管是查询还是修改都是调用其内置的方法

以其中的查询方法为例子

@Override
public <E> List<E> query(Statement statement, ResultHandler resultHandler) throws SQLException {
  PreparedStatement ps = (PreparedStatement) statement;
  ps.execute();
  return resultSetHandler.handleResultSets(ps);
}
复制代码

我们只要通过修改字节码方式将这个类再加载之前修改为如下的方法即可

@Override
public <E> List<E> query(Statement statement, ResultHandler resultHandler) throws SQLException {
long l = System.currentTimeMillis();
  PreparedStatement ps = (PreparedStatement) statement;
  ps.execute();
  return resultSetHandler.handleResultSets(ps);
  System.out.println("执行完毕用时：" + (System.currentTimeMillis() - l));
}
复制代码

代码

新建一个maven项目，pom文件如下

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example</groupId>
    <artifactId>agent_demo</artifactId>
    <version>1.0</version>

    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
    </properties>


    <dependencies>
        <!-- https://mvnrepository.com/artifact/org.javassist/javassist -->
        <dependency>
            <groupId>org.javassist</groupId>
            <artifactId>javassist</artifactId>
            <version>3.27.0-GA</version>
        </dependency>



    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-jar-plugin</artifactId>
                <version>2.4</version>
                <configuration>
                    <archive>
                        <manifest>                         <!--是否要把第三方jar放到manifest的classpath中-->
                            <addClasspath>true</addClasspath>
                        </manifest>
                        <manifestEntries>
                            <Premain-Class>
                                <!-- 入口程序类 -->
                                core.DoMain
                            </Premain-Class>
                        </manifestEntries>
                    </archive>
                </configuration>
            </plugin>
        </plugins>
    </build>

</project>
复制代码

DoMain类如下

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.lang.instrument.Instrumentation;

public class DoMain {

    public static void premain(String args, Instrumentation instrumentation) {
        System.out.println("premain ===> " + args);

        instrumentation.addTransformer(new NotExpiredClassFileTransformer());

    }


    public static void main(String[] args) throws Exception {
        CtClass ctClass = ClassPool.getDefault().get("core.DoMain");
        CtMethod premain = ctClass.getDeclaredMethod("premain", new CtClass[]{
                ClassPool.getDefault().get("java.lang.String"),
                ClassPool.getDefault().get("java.lang.instrument.Instrumentation")
        });

        System.out.println(premain.getLongName());
    }

}
复制代码

修改字节码代码如下

public class SqlTraceClassFileTransformer implements ClassFileTransformer {

    private static final Set<String> CLASS_NAME_SET = new HashSet<>();

    {
        CLASS_NAME_SET.add("org.apache.ibatis.executor.statement.PreparedStatementHandler");
    }


    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (!CLASS_NAME_SET.contains(className)) {
            return null;
        }
        CtClass ctclass = null;
        try {
            ctclass = ClassPool.getDefault().get(className);// 使用全称,用于取得字节码类<使用javassist>

            CtMethod queryMethod = ctclass.getDeclaredMethod("query");
            queryMethod.insertBefore("long l = System.currentTimeMillis();");
            queryMethod.insertAfter("System.out.println("执行完毕用时：" + (System.currentTimeMillis() - l));");

            return ctclass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }

    }

}
复制代码

然后打包成jar

此时在spring boot 项目中加上启动命令

-javaagent:/Users/hans/work_space/java_project/java_agent_demo/target/agent_demo-1.0.jar
复制代码

执行一个sql即可看到控制台打印