自己动手写一个 Arthas 在线诊断工具&系列说明

相关文章：

• 自己动手写一个分库分表中间件（十）线上优化之数据库连接超时优化
• 自己动手写分布式任务调度框架
• 自己动手写 Java 虚拟机（二）-查找 Class 文件
• 自己动手调试 JDK（CLion）
• Java Agent 的简单使用

“自己动手”系列说明

在软件开发的世界中，框架无疑是我们最好的朋友。它们帮助我们抽象复杂性，提供了一种结构化的方式来组织我们的代码，并且通常包含了大量的实用功能，使我们能够更快地开发出高质量的软件。然而，对于许多开发者来说，框架往往是一个黑盒，我们使用它，但往往对其内部的工作原理知之甚少。

系统性地完整学习一个框架是一项耗时且耗力的任务，而且学习的知识很容易被忘记，也很容易走进死胡同。这是因为许多框架经过长时间的迭代，其内部会有各种封装和抽象，这使得开发人员很难快速理解框架的核心逻辑。

在自己动手系列中，将从零开始，实现一些最流行的 Java 框架的核心功能。我的目标是抓大放小，快速理解框架的设计原理和工作机制，从而对框架有一个基本的认知。

通过这种方式，可以为后续更深入地理解这些框架，理解它们为什么会这样设计，以及它们是如何解决复杂问题的学习打下基础。这将帮助我们更好地使用这些框架，也会提高我们的编程技能，使我们能够更好地设计和实现我们自己的代码。

本期“自己动手系列”是实现一个非常 Demo 级别的类 Arthas 诊断工具。可以实现类似于 Arthas watch 命令的功能。

效果

先看一下效果，我先启动一个 Spring Boot 工程，其中有一个 Controller：

/**
 * @author dongguabai
 * @date 2024-01-14 23:22
 */
@RestController
public class TestController {

    @GetMapping("/test")
    public String test(@RequestParam("id") String id) {
        return new Date().toLocaleString() + "->" + id;
    }
}

然后启动自己实现的 Arthas 工具，输入想要诊断的 Spring Boot 应用的进程 ID：

选择 Spring Boot 工程的进程号：

Currently running Java processes：
5600 App
5648 Launcher
5649 Console
5651 Jps
87382 RemoteMavenServer
760 
11934 org.eclipse.equinox.launcher_1.6.700.v20231214-2017.jar
5599 Launcher
> 5600
Attach success
Attach success. Please input the class and method (format: com.example.MyClass#myMethod):
>

再输入想要诊断的函数：

Currently running Java processes：
5600 App
5648 Launcher
5649 Console
5651 Jps
87382 RemoteMavenServer
760 
11934 org.eclipse.equinox.launcher_1.6.700.v20231214-2017.jar
5599 Launcher
> 5600
Attach success
Attach success. Please input the class and method (format: com.example.MyClass#myMethod):
> dongguabai.spring.boot.demo.sb1.TestController#test
Listening on port 56494
Agent applied to dongguabai.spring.boot.demo.sb1.TestController#test
> 

然后调用 dongguabai.spring.boot.demo.sb1.TestController#test 函数：

?  lib curl http://localhost:8080/test\?id\=AS123
2024-1-15 16:26:03->AS123%                                                                                              ?  lib curl http://localhost:8080/test\?id\=AS123
2024-1-15 16:27:06->AS123%                                                                                              ?  lib curl http://localhost:8080/test\?id\=AS123
2024-1-15 16:27:07->AS123%  

再观察诊断工具的控制台：

Currently running Java processes：
5600 App
5648 Launcher
5649 Console
5651 Jps
87382 RemoteMavenServer
760 
11934 org.eclipse.equinox.launcher_1.6.700.v20231214-2017.jar
5599 Launcher
> 5600
Attach success
Attach success. Please input the class and method (format: com.example.MyClass#myMethod):
> dongguabai.spring.boot.demo.sb1.TestController#test
Listening on port 56494
Agent applied to dongguabai.spring.boot.demo.sb1.TestController#test
> Arguments: [AS123]
Return: 2024-1-15 16:26:03->AS123
Arguments: [AS123]
Return: 2024-1-15 16:27:06->AS123
Arguments: [AS123]
Return: 2024-1-15 16:27:07->AS123

可以看到打印出了诊断函数的请求参数和响应参数。

实现原理

控制台

首先需要给用户提供一个交互式的控制台，可以基于 JLine 实现，先看一个简单的 Demo：

public class JLineDemo {

    public static void main(String[] args) throws Exception {
        Terminal terminal = TerminalBuilder.builder().system(true).build();
        LineReader lineReader = LineReaderBuilder.builder().terminal(terminal).build();

        String line;
        while (true) {
            line = lineReader.readLine("> ");
            System.out.println("Input: " + line);
            if ("quit".equalsIgnoreCase(line)) {
                break;
            }
        }
    }
}

运行后会进入一个控制台，可以进行交互：

> 1
Input: 1
> daad
Input: daad
> 

获取正在运行的 JVM 进程

Arthas 启动后会看到当前正在运行的 JVM 进程，可以基于 jps 命令去做这里就涉及到两点：

1. Java 执行 jps 命令
2. 解析命令结果

这也不难，直接使用 Runtime API 即可，再与上面的控制台结合起来：

public class ConsoleDemo {
    public static void main(String[] args) {
        try {
            Terminal terminal = TerminalBuilder.terminal();
            LineReader reader = LineReaderBuilder.builder().terminal(terminal).build();
            Set<String> pids = getRunningJavaProcesses();
            interactWithUser(reader, pids);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static Set<String> getRunningJavaProcesses() throws Exception {
        Set<String> pids = new HashSet<>();
        Process process = Runtime.getRuntime().exec("jps");
        try (BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
            String line;
            System.out.println("Currently running Java processes：");
            while ((line = bufferedReader.readLine()) != null) {
                System.out.println(line);
                String pid = line.split(" ")[0];
                pids.add(pid);
            }
        }
        return pids;
    }

    private static void interactWithUser(LineReader reader, Set<String> pids) {
        while (true) {
            String input = reader.readLine("> ");
            if ("quit".equals(input)) {
                break;
            } else {
                System.out.println("Input：" + input);
            }
        }
    }
}

运行：

Currently running Java processes：
5600 App
87382 RemoteMavenServer
760 
8553 Launcher
8554 ConsoleDemo
8558 Jps
11934 org.eclipse.equinox.launcher_1.6.700.v20231214-2017.jar
5599 Launcher
> 1
Input：1
> 222
Input：222
> 

Java Agent

在Java Agent 的简单使用中介绍过，Java Agent 有两种启动场景，JVM 启动时和运行时候，而目前的场景显然就是在运行时候动态加载 Java Agent。

Attach

可以基于 VirtualMachine 来实现：

private static void attachToProcess(String pid) throws Exception {
    VirtualMachine vm = VirtualMachine.attach(pid);
    System.out.println("Attach success");
    vm.detach();
}

Java Agent

因为是在 JVM 运行时被调用，所以这里要使用 agentmain 函数，同时需要在目标函数执行前后打印出请求参数和响应参数：

public class Agent {

    public static void agentmain(String agentArgs, Instrumentation inst) {
        //initializeAgent
    }
    
    private static void initializeAgent(String agentArgs, final Instrumentation inst) throws Exception {
        String[] args = agentArgs.split("#");
        final String className = args[0];
        final String methodName = args[1];
        System.out.println("className:" + className);
        System.out.println("methodName:" + methodName);
        Class[] allLoadedClasses = inst.getAllLoadedClasses();
        System.out.println(allLoadedClasses.length);

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

                try {
                    ClassPool cp = ClassPool.getDefault();
                    CtClass cc = cp.get(className1.replace("/", "."));
                    CtMethod m = cc.getDeclaredMethod(methodName);
                    System.out.println("Transforming class: " + cc.getName());
                    System.out.println("Transforming method: " + m.getName());
                    m.insertBefore("System.out.println(\"Arguments: \" + java.util.Arrays.toString($args));");
                    m.insertAfter("System.out.println(\"Return: \" + $_);");
                    return cc.toBytecode();
                } catch (Exception e) {
                    System.out.println("Failed to transform class: " + className1);
                    e.printStackTrace();
                    return null;
                }
            }
        }, true);

        for (Class<?> clazz : allLoadedClasses) {
            System.out.println("Loaded class: " + clazz.getName());
            if (clazz.getName().equals(className)) {
                System.out.println("Retransforming class: " + clazz.getName());
                inst.retransformClasses(clazz);
            }
        }
    }
}

pom.xml 如下：

<plugin>
  <groupId>org.apache.maven.plugins</groupId>
  <artifactId>maven-jar-plugin</artifactId>
  <version>3.2.0</version>
  <configuration>
    <archive>
      <manifest>
        <addClasspath>true</addClasspath>
        <mainClass>blog.dongguabai.arthas.Agent</mainClass>
        <addDefaultImplementationEntries>true</addDefaultImplementationEntries>
      </manifest>
      <manifestEntries>
        <Agent-Class>blog.dongguabai.arthas.Agent</Agent-Class>
        <Can-Retransform-Classes>true</Can-Retransform-Classes>
      </manifestEntries>
    </archive>
  </configuration>
</plugin>

IPC 机制

这里有个细节，控制台（也就是 Arthas.jar）是一个独立的 JVM 进程，Java Agent 从目标 JVM 进程采集到数据后如何返回给控制台呢。其实也很简单，通过 Socket 通信即可。

Agent 直接通过 Socket 发送数据给控制台：

 m.insertBefore("{ java.net.Socket socket = new java.net.Socket(\"localhost\", " + port + "); " +
                        "java.io.PrintWriter out = new java.io.PrintWriter(socket.getOutputStream(), true); " +
                        "out.println(\"Arguments: \" + java.util.Arrays.toString($args)); " +
                        "out.close(); socket.close(); }");
                m.insertAfter("{ java.net.Socket socket = new java.net.Socket(\"localhost\", " + port + "); " +
                        "java.io.PrintWriter out = new java.io.PrintWriter(socket.getOutputStream(), true); " +
                        "out.println(\"Return: \" + $_); " +
                        "out.close(); socket.close(); }");

控制台接收 Socket 数据并且打印：

    private static void applyAgent(String pid, String className, String methodName) throws Exception {
        VirtualMachine vm = VirtualMachine.attach(pid);
        ServerSocket serverSocket = new ServerSocket(0);
        int port = serverSocket.getLocalPort();
        new Thread(() -> {
            System.out.println("Listening on port " + port);
            while (true) {
                try (Socket clientSocket = serverSocket.accept();
                     BufferedReader in = new BufferedReader(new InputStreamReader(clientSocket.getInputStream()))) {
                    String inputLine;
                    while ((inputLine = in.readLine()) != null) {
                        System.out.println(inputLine);
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }).start();
        ...
    }

总结

本文实现一个非常 Demo 级别的类 Arthas 诊断工具，其中涉及到了 JLine、Java Agent、ClassLoader、Javassist、Socket 等技术。

尽管工具还很简单，功能也很有限，但它为理解和学习如 Arthas 这样的复杂诊断工具提供了一个实践入口。后续也会继续扩展这个工具，添加更多的功能，如更详细的方法追踪，更丰富的 JVM 信息获取，甚至是内存和 CPU 的监控等。
源码地址：https://gitee.com/dongguabai/blog/tree/master/dongguabai-arthas

欢迎关注公众号：