优秀的编程知识分享平台

网站首页 > 技术文章 正文

Java agent结合javassist为方法增加耗时统计

nanyue 2024-09-21 20:06:13 技术文章 5 ℃

java agent介绍

向我们熟知的spring boot devtool、Jbrebel、某些破解程序都是使用Java agent来实现的。

常见的使用方式为指定JVM参数-javaagent=/path/to/agent.jar=param。

java agent分为main之前执行的premain和main之后执行的agentmain。

main之前的agent

public class MyAgentV2 {

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

然后在agent jar的META-INF/MANIFEST.MF中增加Premain-Class=MyAgentV2的全限定名。

如果使用maven,maven-jar-plugin中配置为

<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-jar-plugin</artifactId>
<version>3.2.0</version>
<configuration>
    <archive>
        <manifestEntries>
            <Project-name>${project.name}</Project-name>
            <Project-version>${project.version}</Project-version>
            <Premain-Class>com.donny.agent.MyAgent</Premain-Class>
        </manifestEntries>
    </archive>
    <skip>true</skip>
</configuration>
</plugin>

使用Agent:在调用方加上jvm参数-javaagent=/path/to/agent.jar

main之后的agent

public class TestAgent {
  public static void agentmain(String args, Instrumentation instrumentation) {
    System.out.println("agent run after main.");
  }
}

在agent.jar的META-INF/MANIFEST.MF中增加Agent-Class=TestAgent的全限定名。

使用Agent:在调用方加上jvm参数-javaagent=/path/to/agent.jar

使用Instrumentation修改Class字节码

如果使用premain agent模式,在类被ClassLoader加载之前会进入transform方法。可以通过它来修改Class字节码。

public class MyAgentV2 {

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

    instrumentation.addTransformer(
        new ClassFileTransformer() {
          public byte[] transform(
              ClassLoader loader,
              String className,
              Class<?> classBeingRedefined,
              ProtectionDomain protectionDomain,
              byte[] classfileBuffer)
              throws IllegalClassFormatException {
                // 拦截HelloService,让HelloService使用老版本的HelloService
              if (!"com/donny/agent/HelloService".equals(className)) {
                  // return null表示使用原Class文件的内容
                  return null;
              }

              // 读取HelloServiceV1的HelloService的字节码,即/target/HelloService.class
              InputStream in = getClass().getResourceAsStream("/HelloService.class");
              try {
                  return IOUtils.toByteArray(in);
              } catch (IOException e) {
                  e.printStackTrace();
              }
              return null;
          }
        });
  }
}

在上面的示例中,我们将com/donny/agent/HelloService的字节码替换为了/target/HelloService.class的字节码。

注意:如果在执行transform之前类已经被ClassLoader加载了,则不会进入transform。

已被加载的Class如何重新进入transform?

可以使用Instrumetation提供的retransformClasses方法,同时需要设置META-INF/MANIFEST.MF中Can-Retransform-Classes=true。

通过retransformClasses方法,Class会从ClassLoader中清除,下次使用该Class时会重新进入transform。

public class MyAgentV4 {

  public static void premain(String arg, Instrumentation instrumentation)
      throws IOException, UnmodifiableClassException, ClassNotFoundException {
    System.out.println("premain,arg = " + arg);

    HelloService hs = new HelloService();
    // 这里调不掉用方法都无所谓
     hs.sayHello("jack",100); // 打印HelloServiceV2

    instrumentation.addTransformer(
        new ClassFileTransformer() {
          public byte[] transform(
              ClassLoader loader,
              String className,
              Class<?> classBeingRedefined,
              ProtectionDomain protectionDomain,
              byte[] classfileBuffer)
              throws IllegalClassFormatException {
            // 拦截HelloService,让HelloService使用老版本的HelloService
            if (!"com/donny/agent/HelloService".equals(className)) {
              // return null表示使用原Class文件的内容
              return null;
            }
            InputStream in = MyAgentV4.class.getResourceAsStream("/HelloService.class");
            try {
              return IOUtils.toByteArray(in);
            } catch (IOException e) {
              e.printStackTrace();
            }
            ;
            return null;
          }
        },
        true);

    // 重置HelloService,重置后,会再次进入transform流程。
    instrumentation.retransformClasses(HelloService.class);

    hs.sayHello("jack ma",99); // 打印HelloServiceV1.0
  }
}

不重置Class,是否可以修改某个Class?

可以通过Instrumentation的redefineClasses方法,并配置META-INF/MANIFEST.MF中Can-Redefine-Classes=true来重新定义该Class。

public class MyAgentV3 {

  public static void premain(String arg, Instrumentation instrumentation)
      throws IOException, UnmodifiableClassException, ClassNotFoundException {
    System.out.println("premain,arg = " + arg);

    HelloService hs = new HelloService();
    // 这里调不掉用方法都无所谓
    // hs.sayHello("jack");

    // 重新定义HelloService
    InputStream in = MyAgentV3.class.getResourceAsStream("/HelloService.class");
    byte[] bytes = IOUtils.toByteArray(in);
    instrumentation.redefineClasses(new ClassDefinition(HelloService.class, bytes));
  }
}

集合javassist来实现对指定匹配模式的包下的所有方法增加调用耗时统计

javassist简介

javassist是一个用来处理Java字节码的库,可以用来在编译好的类中新增或修改方法或类。它是jboss下的一个子项目,主要优点在于简单、快速。

在本示例中,我们用来它新增和重命名方法。

javassist特殊语法

  • $0,$1,$2,...
    方法参数表示,1表示第1个参数,表示第二个参数,以此类推。静态方法没有,所以0不存在。
  • $args
    将参数以Object数组的形式进行封装,相当于new Object[]{2,...}
  • $cflow(...)
    方法在递归调用时,可读取其递归调用的层次
  • $r
    用于封装方法结果,如:Object result = ...;return ($r)result;
  • $w
    用于将基础类型转换为包装类型,如Integer a = ($w)123;
  • $_
    设置返回结果w)1;相当于return ($w)1;
  • $sig
    获取方法中所有参数类型,数组形式展现
  • $type
    获取方法结果的Class
  • $class
    获取当前方法所在类的Class。
  • $
    自动填充方法参数。

实现思路

1.根据指定的包扫描包下的所有类; 2.在transform方法中,对扫描到的类的所有方法通过javassist动态修改,实现调用耗时的统计;

需要注意的是:虽然javassist提供了insertBefore/After/At方法,但实际每次执行它有自己的作用域,外部无法访问。

比如在方法执行前打印当前时间

ClassPool classPool = new ClassPool();
classPool.appendSystemPath();
CtClass ctClass = classPool.get("com.donny.agent.HelloService");
CtMethod ctMethod = ctClass.getDeclaredMethod("sayHello");
// 打印当前时间,这里实际效果为{System.out.println(System.currentTimeMillis());},会被{}包裹起来,外部无法访问。
ctMethod.insertBefore("System.out.println(System.currentTimeMillis());");

所以我们的实现思路是新增一个方法,在新增的方法中调用原来的方法,并加入耗时统计的代码。 原方法定义如下:

public void sayHello(String name) {
   System.out.println(name);
}

我们需要改造为下面的代码:

public void sayHello$agent(String name) {
   System.out.println(name);
}

public void sayHello(String name) {
    long start = System.currentTimeMillis();
    try {
        sayHello$agent(name);
    }finally {
        System.out.println(System.currentTimeMillis() - start);
    }   
}

实现代码

Class扫描的工具类:

/** 根据报名扫描包/子包下的所有Class */
public final class ClassUtils {
  public static Set<String> scanClassesByPackage(String packageName)
      throws IOException, ClassNotFoundException {
    Set<String> classes = new HashSet<>();
    if (packageName == null || packageName.isEmpty()) {
      return classes;
    }

    String pkgName = packageName;
    packageName = packageName.replaceAll("\\.", "/");
    Enumeration<URL> urls = MyAgent.class.getClassLoader().getResources(packageName);
    while (urls.hasMoreElements()) {
      URL url = urls.nextElement();
      String protocol = url.getProtocol();
      if (protocol.equals("file")) {
        scanClassInPackageByFile(pkgName, url.getPath(), classes);
      } else if (protocol.equals("jar")) {
        JarFile jarFile = ((JarURLConnection) url.openConnection()).getJarFile();
        scanClassesInPackageByJar(pkgName, jarFile, classes);
      }
    }
    return classes;
  }

  private static void scanClassesInPackageByJar(
      String packageName, JarFile jarFile, Set<String> classes) {
    //    Enumeration<JarEntry> entries = jarFile.entries();
    //    while (entries.hasMoreElements()) {
    //      JarEntry entry = entries.nextElement();
    //
    //    }
    // todo
  }

  private static void scanClassInPackageByFile(
      String basePackageName, String path, Set<String> classes) throws ClassNotFoundException {
    File file = new File(path);
    File[] files =
        file.listFiles(
            new FileFilter() {

              @Override
              public boolean accept(File f) {
                return f.isDirectory() || f.getName().endsWith(".class");
              }
            });

    if (files == null) return;

    for (File ff : files) {
      if (ff.isDirectory()) {
        scanClassInPackageByFile(
            basePackageName + "." + ff.getName(), ff.getAbsolutePath(), classes);
      } else {
        String classname =
            basePackageName + "." + ff.getName().substring(0, ff.getName().length() - 6);
        //Class<?> clazz = Thread.currentThread().getContextClassLoader().loadClass(classname);
        classes.add(classname);
      }
    }
  }
}

Agent类:

public class MyAgent {

  public static void premain(String args, Instrumentation instrumentation) throws IOException, ClassNotFoundException, NotFoundException, CannotCompileException {
    System.out.println("MyAgent is running.args = " + args);
    // 这里可以按实际需要修改匹配模式,比如实现com.donny.*.*Service
    String pattern = "com.donny.agent.*";
    String packageName = pattern.substring(0,pattern.indexOf(".*"));
    // 扫描通配符配置的Class
    final Set<String> classes = ClassUtils.scanClassesByPackage(packageName);
    ClassPool cp = new ClassPool();
    cp.appendSystemPath();

    instrumentation.addTransformer(new ClassFileTransformer() {
      @Override
      public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] bytes = null;
        String realClassName = className.replace("/", ".");
        if (!classes.contains(realClassName)) {
          return bytes;
        }

        try {
          CtClass cc = cp.get(realClassName);
          bytes = buildMonitor(cc);
        } catch (NotFoundException e) {
          e.printStackTrace();
        } catch (CannotCompileException e) {
          e.printStackTrace();
        } catch (IOException e) {
          e.printStackTrace();
        }
        return bytes;
      }
    });
  }

  public static TraceInfo begin(Object[] args) {
    return new TraceInfo(System.currentTimeMillis(), args);
  }

  public static void end(TraceInfo traceInfo) {
    System.out.println("执行时间:" + (System.currentTimeMillis() - traceInfo.begin));
    System.out.println("方法参数 = " + Arrays.toString(traceInfo.params));
  }

  public static class TraceInfo {
    private long begin;
    private Object[] params;

    public TraceInfo() {}

    public TraceInfo(long begin, Object[] params) {
      this.begin = begin;
      this.params = params;
    }
  }

  private static byte[] buildMonitor(CtClass ctClass) throws NotFoundException, CannotCompileException, IOException {
    boolean hasMethod = false;
    for (CtMethod me : ctClass.getDeclaredMethods()) {
      // 只对public方法进行插桩
      if (AccessFlag.isPublic(me.getModifiers())) {
        // 排除main方法
        // if (me.getName().equals("main")) continue;
        hasMethod = true;
        // 对原有方法进行copy
        CtMethod copy = CtNewMethod.copy(me, ctClass, new ClassMap());
        // 将原有方法重命名
        String agentName = me.getName() + "$agent";
        me.setName(agentName);
        // 调用原方法
        String callerAgent = String.format("return %s($);", agentName);
        // 在copy的方法中调用原有方法,并加上监控数据
        copy.setBody(
            "{\n"
                + "        com.donny.agent.MyAgent.TraceInfo traceInfo = com.donny.agent.MyAgent.begin($args);\n"
                + "        try {\n"
                + callerAgent
                + "        } finally {\n"
                + "            com.donny.agent.MyAgent.end(traceInfo);\n"
                + "        }\n"
                + "    }");
        // 添加拷贝的方法到Class
        ctClass.addMethod(copy);
      }
    }

    return !hasMethod ? null : ctClass.toBytecode();
  }
}

待测试的Service:

public class HelloService {

  public int sayHello(String name,int age) {
    System.out.println("HelloServiceV2,name = " + name + ",age = " + age);
    return 1;
  }

  public void testHeavyBusiness() {
    try {
      Thread.sleep(1000);
    } catch (InterruptedException e) {
      e.printStackTrace();
    }
    System.out.println("testHeavyBusiness执行完毕。");
  }
}

测试类:

public class Test {
  public static void main(String[] args) {
    HelloService hs = new HelloService();
    int result = hs.sayHello("jack",129);
    System.out.println("result = " + result);

    hs.testHeavyBusiness();
  }
}

测试步骤:

  1. 将agent打为jar包(参考上面的内容);
  2. 编辑Test配置,增加jvm参数:-javaagent=/path/to/agent.jar
  3. 运行Test.java 控制台输出:
HelloServiceV2,name = jack,age = 129
执行时间:0
方法参数 = [jack, 129]
result = 1
testHeavyBusiness执行完毕。
执行时间:1000
方法参数 = []

结语

本文简单介绍了java agent技术结合javassist动态修改字节码技术来实现动态为方法添加调用耗时统计的功能。这块在热加载、监控等方面使用较多,一般很少用到,算是作为一个了解。如果有收集监控数据的需求,可以深入了解java agent相关的东西。

最近发表
标签列表