用javaagent和javassist实现Arthas的watch功能

news2024/9/27 21:22:53

一、被监控的服务 spring-boot-demo

1、 pom.xml

<?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>com.techhf</groupId>
    <artifactId>spring-boot-demo</artifactId>
    <version>1.0-SNAPSHOT</version>
    
    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
            <version>2.3.4.RELEASE</version>
        </dependency>
    
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-undertow</artifactId>
            <version>2.3.4.RELEASE</version>
        </dependency>
    </dependencies>
    
</project>

2、DemoController

package com.techhf.demo.controller;

import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author shiweijia
 * @date 2024/8/1 17:31
 */
@RequestMapping("/demo")
@RestController
public class DemoController {
    
    @RequestMapping("/hello")
    public String hello(String name) {
        return "hello world " + name;
    }
    
    
}

3、DemoApplication

package com.techhf.demo;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

/**
 * @author shiweijia
 * @date 2024/8/1 17:30
 */
@SpringBootApplication
public class DemoApplication {
    
    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
    }
}

4、运行DemoApplication

-javaagent:D:/temp/jvm-agent-1.0-SNAPSHOT-jar-with-dependencies.jar

二、jvm-agent

1、pom.xml

<?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>com.techhf</groupId>
    <artifactId>jvm-agent</artifactId>
    <version>1.0-SNAPSHOT</version>
    
    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
    </properties>
    
    <dependencies>
        <dependency>
            <groupId>org.javassist</groupId>
            <artifactId>javassist</artifactId>
            <version>3.30.2-GA</version>
        </dependency>
    </dependencies>
    
    
    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-assembly-plugin</artifactId>
                <version>3.3.0</version>
                <configuration>
                    <descriptorRefs>
                        <descriptorRef>jar-with-dependencies</descriptorRef>
                    </descriptorRefs>
                    <archive>
                        <manifestEntries>
                            <Premain-Class>com.techhf.agent.WatchAgent</Premain-Class>
                            <Agent-Class>com.techhf.agent.WatchAgent</Agent-Class>
                            <Can-Redefine-Classes>true</Can-Redefine-Classes>
                            <Can-Retransform-Classes>true</Can-Retransform-Classes>
                        </manifestEntries>
                    </archive>
                </configuration>
                <executions>
                    <execution>
                        <id>make-assembly</id>
                        <phase>package</phase>
                        <goals>
                            <goal>single</goal>
                        </goals>
                    </execution>
                </executions>
            </plugin>
        </plugins>
    </build>

</project>

2、WatchAgent

package com.techhf.agent;

import java.lang.instrument.Instrumentation;

/**
 * @author shiweijia
 * @date 2024/8/1 16:00
 */
public class WatchAgent {
    
    public static void premain(String agentArgs, Instrumentation inst) {
        agentArgs = "com.techhf.demo.controller.DemoController,hello";
        System.out.println("premain attached with args: " + agentArgs);
        transform(agentArgs, inst);
    }
    
    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("agentmain attached with args: " + agentArgs);
        transform(agentArgs, inst);
    }
    
    private static void transform(String agentArgs, Instrumentation inst) {
        // 解析 agentArgs 来获取类名和方法名
        String[] parts = agentArgs.split(",");
        String className = parts[0];
        String methodName = parts[1];
        inst.addTransformer(new WatchTransformer(className, methodName));
    }
    
}

3、WatchTransformer

package com.techhf.agent;

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

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @author shiweijia
 * @date 2024/8/3 16:56
 */
public class WatchTransformer implements ClassFileTransformer {
    
    private String className;
    
    private String methodName;
    
    public WatchTransformer(String className, String methodName) {
        this.className = className;
        this.methodName = methodName;
    }
    
    
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace('/', '.');
        if (this.className.equals(className)) {
            try {
                ClassPool pool = ClassPool.getDefault();
                CtClass cc = pool.get(className);
                
                CtMethod method = cc.getDeclaredMethod(methodName);
                String beforeCode = "{ System.out.print(\"Before Method Name: \" + \"" + method.getName() + "\" + \", Parameters: \"); " +
                        "for (int i = 0; i < $args.length; i++) { " +
                        "    System.out.print($args[i]); " +
                        "} " +
                        "System.out.println(); }";
                method.insertBefore(beforeCode);
                
                method.insertAfter("{ System.out.println(\"After Method Name: \"+\"" + method.getName()+ "\" + \", Return: \" + $_); }");
                
                return cc.toBytecode();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return classfileBuffer;
    }
}

三、启动时监控

如上述实例，运行时加入-javaagent相关参数即可。

四、运行时监控

import com.sun.tools.attach.VirtualMachine;

/**
 * @author shiweijia
 * @date 2024/8/1 17:35
 */
public class AttachLoader {
    
    public static void main(String[] args) {
        if (args.length < 2) {
            System.out.println("Usage: AttachLoader <pid> <agent-path> <agent-args>");
            System.exit(1);
        }
        String pid = args[0];
        String agentPath = args[1];
        String agentArgs = args[2];
        try {
            // 获取目标 JVM 进程
            VirtualMachine vm = VirtualMachine.attach(pid);
            vm.loadAgent(agentPath, agentArgs);
            // 可选：在需要时分离（detach）
            // vm.detach();
            System.out.println("Agent loaded into PID " + pid);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    
}

运行参数为：