BTrace实现浅析

在之前的文章中我们简单介绍了BTrace的用法，今天我们通过源代码来看看BTrace是如何实现的。

从BTrace的启动脚本中可以找到相关入口，

${JAVA_HOME}/bin/java -Dcom.sun.btrace.probeDescPath=. -Dcom.sun.btrace.dumpClasses=false -Dcom.sun.btrace.debug=false -Dcom.sun.btrace.unsafe=false -cp ${BTRACE_HOME}/build/btrace-client.jar:${TOOLS_JAR}:/usr/share/lib/java/dtrace.jar com.sun.btrace.client.Main $*

所以Main-Class就是com.sun.btrace.client.Main了，来看看它的main方法，

   public static void main(String[] args) {
        ////////////////////////////////////////
        // 1. 参数解析
        ////////////////////////////////////////
        if (args.length < 2) {
            usage();
        }

        int port = BTRACE_DEFAULT_PORT;
        String classPath = ".";
        String includePath = null;

        int count = 0;
        boolean portDefined = false;
        boolean classpathDefined = false;
        boolean includePathDefined = false;

        for (;;) {
            if (args[count].charAt(0) == '-') {
                if (args.length <= count+1) {
                    usage();
                }
                if (args[count].equals("-p") && !portDefined) {
                    try {
                        port = Integer.parseInt(args[++count]);
                        if (isDebug()) debugPrint("accepting port " + port);
                    } catch (NumberFormatException nfe) {
                        usage();
                    }
                    portDefined = true;
                } else if ((args[count].equals("-cp") ||
                    args[count].equals("-classpath"))
                    && !classpathDefined) {
                    classPath = args[++count];
                    if (isDebug()) debugPrint("accepting classpath " + classPath);
                    classpathDefined = true;
                } else if (args[count].equals("-I") && !includePathDefined) {
                    includePath = args[++count];
                    if (isDebug()) debugPrint("accepting include path " + includePath);
                    includePathDefined = true;
                } else {
                    usage();
                }
                count++;
                if (count >= args.length) {
                    break;
                }
            } else {
                break;
            }
        }

        if (! portDefined) {
            if (isDebug()) debugPrint("assuming default port " + port);
        }

        if (! classpathDefined) {
            if (isDebug()) debugPrint("assuming default classpath '" + classPath + "'");
        }

        if (args.length < (count + 1)) {
            usage();
        }

        String pid = args[count];
        String fileName = args[count + 1];
        String[] btraceArgs = new String[args.length - count];
        if (btraceArgs.length > 0) {
            System.arraycopy(args, count, btraceArgs, 0, btraceArgs.length);
        }

        try {
            Client client = new Client(port, PROBE_DESC_PATH,
                DEBUG, TRACK_RETRANSFORM, UNSAFE, DUMP_CLASSES, DUMP_DIR);
            if (! new File(fileName).exists()) {
                errorExit("File not found: " + fileName, 1);
            }

            ////////////////////////////////////////
            // 2. 编译btrace脚本
            ////////////////////////////////////////
            byte[] code = client.compile(fileName, classPath, includePath);
            if (code == null) {
                errorExit("BTrace compilation failed", 1);
            }

            ////////////////////////////////////////
            // 3. attach到目标VM
            ////////////////////////////////////////
            client.attach(pid);
            registerExitHook(client);
            if (con != null) {
                registerSignalHandler(client);
            }
            if (isDebug()) debugPrint("submitting the BTrace program");

            ////////////////////////////////////////
            // 4. 提交btrace请求
            ////////////////////////////////////////
            client.submit(fileName, code, btraceArgs,
                createCommandListener(client));
        } catch (IOException exp) {
            errorExit(exp.getMessage(), 1);
        }
    }

BTrace脚本的编译细节（包括脚本解析等）我们暂不深究。看第3步之前先来看第4步提交的请求，com.sun.btrace.client.Client#submit，

    /**
     * Submits the compiled BTrace .class to the VM
     * attached and passes given command line arguments.
     * Receives commands from the traced JVM and sends those
     * to the command listener provided.
     */
    public void submit(String fileName, byte[] code, String[] args,
            CommandListener listener) throws IOException {
        if (sock != null) {
            throw new IllegalStateException();
        }
        submitDTrace(fileName, code, args, listener);
        try {
            if (debug) {
                debugPrint("opening socket to " + port);
            }
            ////////////////////////////////////////
            // 与目标VM通过Socket进行通信
            ////////////////////////////////////////
            sock = new Socket("localhost", port);
            oos = new ObjectOutputStream(sock.getOutputStream());
            if (debug) {
                debugPrint("sending instrument command");
            }
            ////////////////////////////////////////
            // 给目标VM发送InstrumentCommand
            ////////////////////////////////////////
            WireIO.write(oos, new InstrumentCommand(code, args));
            ois = new ObjectInputStream(sock.getInputStream());
            if (debug) {
                debugPrint("entering into command loop");
            }
            commandLoop(listener);
        } catch (UnknownHostException uhe) {
            throw new IOException(uhe);
        }
    }

现在我们再来看第3步，也就是com.sun.btrace.client.Client#attach，

    /**
     * Attach the BTrace client to the given Java process.
     * Loads BTrace agent on the target process if not loaded
     * already.
     */
    public void attach(String pid) throws IOException {
        try {
            String agentPath = "/btrace-agent.jar";
            String tmp = Client.class.getClassLoader().getResource("com/sun/btrace").toString();
            tmp = tmp.substring(0, tmp.indexOf("!"));
            tmp = tmp.substring("jar:".length(), tmp.lastIndexOf("/"));
            agentPath = tmp + agentPath;
            agentPath = new File(new URI(agentPath)).getAbsolutePath();
            attach(pid, agentPath, null, null);
        } catch (RuntimeException re) {
            throw re;
        } catch (IOException ioexp) {
            throw ioexp;
        } catch (Exception exp) {
            throw new IOException(exp.getMessage());
        }
    }

    /**
     * Attach the BTrace client to the given Java process.
     * Loads BTrace agent on the target process if not loaded
     * already. Accepts the full path of the btrace agent jar.
     * Also, accepts system classpath and boot classpath optionally.
     */
    public void attach(String pid, String agentPath, String sysCp, String bootCp) throws IOException {
        try {
            VirtualMachine vm = null;
            if (debug) {
                debugPrint("attaching to " + pid);
            }
            vm = VirtualMachine.attach(pid);
            if (debug) {
                debugPrint("checking port availability: " + port);
            }
            Properties serverVmProps = vm.getSystemProperties();
            int serverPort = Integer.parseInt(serverVmProps.getProperty("btrace.port", "-1"));
            if (serverPort != -1) {
                if (serverPort != port) {
                    throw new IOException("Can not attach to PID " + pid + " on port " + port + ". There is already a BTrace server active on port " + serverPort + "!");
                }
            } else {
                if (!isPortAvailable(port)) {
                    throw new IOException("Port " + port + " unavailable.");
                }
            }

            if (debug) {
                debugPrint("attached to " + pid);
            }

            if (debug) {
                debugPrint("loading " + agentPath);
            }
            String agentArgs = "port=" + port;
            if (debug) {
                agentArgs += ",debug=true";
            }
            if (unsafe) {
                agentArgs += ",unsafe=true";
            }
            if (dumpClasses) {
                agentArgs += ",dumpClasses=true";
                agentArgs += ",dumpDir=" + dumpDir;
            }
            if (trackRetransforms) {
                agentArgs += ",trackRetransforms=true";
            }
            if (bootCp != null) {
                agentArgs += ",bootClassPath=" + bootCp;
            }
            if (sysCp == null) {
                sysCp = getToolsJarPath(
                    serverVmProps.getProperty("java.class.path"),
                    serverVmProps.getProperty("java.home")
                );
            }
            String cmdQueueLimit = System.getProperty(BTraceRuntime.CMD_QUEUE_LIMIT_KEY, null);
            if (cmdQueueLimit != null) {
                agentArgs += ",cmdQueueLimit=" + cmdQueueLimit;
            }
            agentArgs += ",systemClassPath=" + sysCp;
            agentArgs += ",probeDescPath=" + probeDescPath;
            if (debug) {
                debugPrint("agent args: " + agentArgs);
            }
            vm.loadAgent(agentPath, agentArgs);
            if (debug) {
                debugPrint("loaded " + agentPath);
            }
        } catch (RuntimeException re) {
            throw re;
        } catch (IOException ioexp) {
            throw ioexp;
        } catch (Exception exp) {
            throw new IOException(exp.getMessage());
        }
    }

可以看到这个地方使用了Attach API。最后调用了VirtualMachine#loadAgent方法，加载的agent是$BTRACE_HOME/build/btrace-agent.jar，它的MANIFEST.MF是这样的，

Manifest-Version: 1.0
Ant-Version: Apache Ant 1.8.0
Created-By: 1.7.0_07-b10 (Oracle Corporation)
Premain-Class: com.sun.btrace.agent.Main
Agent-Class: com.sun.btrace.agent.Main
Boot-Class-Path: btrace-boot.jar
Can-Redefine-Classes: true
Can-Retransform-Classes: true

VirtualMachine#loadAgent的时候会调用Agent-Class的agentmain方法，这里也就是com.sun.btrace.agent.Main#agentmain，

   public static void agentmain(String args, Instrumentation inst) {
        main(args, inst);
    }

   private static synchronized void main(final String args, final Instrumentation inst) {
        if (Main.inst != null) {
            return;
        } else {
            Main.inst = inst;
        }

        ////////////////////////////////////////
        // 1. 参数解析
        ////////////////////////////////////////
        if (isDebug()) debugPrint("parsing command line arguments");
        parseArgs(args);
        if (isDebug()) debugPrint("parsed command line arguments");

        /////// Boot-Class-Path: btrace-boot.jar
        String bootClassPath = argMap.get("bootClassPath");
        if (bootClassPath != null) {
            if (isDebug()) {
                 debugPrint("Bootstrap ClassPath: " + bootClassPath);
            }
            StringTokenizer tokenizer = new StringTokenizer(bootClassPath, File.pathSeparator);
            try {
                while (tokenizer.hasMoreTokens()) {
                    String path = tokenizer.nextToken();
                    inst.appendToBootstrapClassLoaderSearch(new JarFile(new File(path)));
                }
            } catch (IOException ex) {
                debugPrint("adding to boot classpath failed!");
                debugPrint(ex);
                return;
            }
        }

        String systemClassPath = argMap.get("systemClassPath");
        if (systemClassPath != null) {
            if (isDebug()) {
                 debugPrint("System ClassPath: " + systemClassPath);
            }
            StringTokenizer tokenizer = new StringTokenizer(systemClassPath, File.pathSeparator);
            try {
                while (tokenizer.hasMoreTokens()) {
                    String path = tokenizer.nextToken();
                    inst.appendToSystemClassLoaderSearch(new JarFile(new File(path)));
                }
            } catch (IOException ex) {
                debugPrint("adding to boot classpath failed!");
                debugPrint(ex);
                return;
            }
        }

        String tmp = argMap.get("noServer");
        boolean noServer = tmp != null && !"false".equals(tmp);
        if (noServer) {
            if (isDebug()) debugPrint("noServer is true, server not started");
            return;
        }

        ////////////////////////////////////////
        // 2. 启动agent线程
        ////////////////////////////////////////
        Thread agentThread = new Thread(new Runnable() {
            public void run() {
                BTraceRuntime.enter();
                try {
                    startServer();
                } finally {
                    BTraceRuntime.leave();
                }
            }
        });
        BTraceRuntime.enter();
        try {
            agentThread.setDaemon(true);
            if (isDebug()) debugPrint("starting agent thread");
            agentThread.start();
        } finally {
            BTraceRuntime.leave();
        }
    }

startServer方法，

  private static void startServer() {
        int port = BTRACE_DEFAULT_PORT;
        String p = argMap.get("port");
        if (p != null) {
            try {
                port = Integer.parseInt(p);
            } catch (NumberFormatException exp) {
                error("invalid port assuming default..");
            }
        }
        ServerSocket ss;
        try {
            if (isDebug()) debugPrint("starting server at " + port);
            System.setProperty("btrace.port", String.valueOf(port));
            if (scriptOutputFile != null && scriptOutputFile.length() > 0) {
                System.setProperty("btrace.output", scriptOutputFile);
            }
            ss = new ServerSocket(port);
        } catch (IOException ioexp) {
            ioexp.printStackTrace();
            return;
        }

        while (true) {
            try {
                if (isDebug()) debugPrint("waiting for clients");
                ////////////////////////////////////////
                // 等待客户端连接上来
                ////////////////////////////////////////
                Socket sock = ss.accept();
                if (isDebug()) debugPrint("client accepted " + sock);
                ////////////////////////////////////////
                // 生成RemoteClient
                ////////////////////////////////////////
                Client client = new RemoteClient(inst, sock);
                registerExitHook(client);
                ////////////////////////////////////////
                // 处理客户端请求
                ////////////////////////////////////////
                handleNewClient(client);
            } catch (RuntimeException re) {
                if (isDebug()) debugPrint(re);
            } catch (IOException ioexp) {
                if (isDebug()) debugPrint(ioexp);
            }
        }
    }

来看下RemoteClient的构造函数，

   RemoteClient(Instrumentation inst, Socket sock) throws IOException {
        super(inst);
        this.sock = sock;
        this.ois = new ObjectInputStream(sock.getInputStream());
        this.oos = new ObjectOutputStream(sock.getOutputStream());
        ////////////////////////////////////////
        // 读取客户端提交过来的InstrumentCommand
        ////////////////////////////////////////
        Command cmd = WireIO.read(ois);
        if (cmd.getType() == Command.INSTRUMENT) {
            if (debug) Main.debugPrint("got instrument command");
            ////////////////////////////////////////
            // 保存编译后的btrace脚本代码到Client#btraceCode
            ////////////////////////////////////////
            Class btraceClazz = loadClass((InstrumentCommand)cmd);
            if (btraceClazz == null) {
                throw new RuntimeException("can not load BTrace class");
            }
        } else {
            errorExit(new IllegalArgumentException("expecting instrument command!"));
            throw new IOException("expecting instrument command!");
        }
        ...
    }

处理客户端请求，

  private static void handleNewClient(final Client client) {
        serializedExecutor.submit(new Runnable() {

            public void run() {
                try {
                    if (isDebug()) debugPrint("new Client created " + client);
                    if (client.shouldAddTransformer()) {
                        /////////////////////////////////
                        // 1. 添加ClassFileTransformer
                        /////////////////////////////////
                        client.registerTransformer();

                        /////////////////////////////////
                        // 2. 获取满足脚本中条件的全部类
                        /////////////////////////////////
                        Class[] classes = inst.getAllLoadedClasses();
                        ArrayList<Class> list = new ArrayList<Class>();
                        if (isDebug()) debugPrint("filtering loaded classes");
                        for (Class c : classes) {
                            if (inst.isModifiableClass(c) && client.isCandidate(c)) {
                                if (isDebug()) debugPrint("candidate " + c + " added");
                                list.add(c);
                            }
                        }
                        list.trimToSize();
                        int size = list.size();
                        if (isDebug()) debugPrint("added as ClassFileTransformer");
                        if (size > 0) {
                            classes = new Class[size];
                            list.toArray(classes);
                            client.startRetransformClasses(size);
                            /////////////////////////////////
                            // 3. 开始进行retransform
                            /////////////////////////////////
                            if (isDebug()) {
                                for(Class c : classes) {
                                    try {
                                      inst.retransformClasses(c);
                                    } catch (VerifyError e) {
                                        debugPrint("verification error: " + c.getName());
                                    }
                                }
                            } else {
                                inst.retransformClasses(classes);
                            }
                            client.skipRetransforms();
                        }
                    }
                    client.getRuntime().send(new OkayCommand());
                } catch (UnmodifiableClassException uce) {
                    if (isDebug()) {
                        debugPrint(uce);
                    }
                    client.getRuntime().send(new ErrorCommand(uce));
                }
            }
        });

    }

com.sun.btrace.agent.Client#registerTransformer方法中会调用java.lang.instrument.Instrumentation#addTransformer，

      void registerTransformer() {
            inst.addTransformer(clInitTransformer, false);
            inst.addTransformer(this, true);
        }

其实主要就是Attach API的使用，通过java.lang.instrument.Instrumentation#addTransformer添加了ClassFileTransformer，当调用java.lang.instrument.Instrumentation#retransformClasses时，上面所添加的ClassFileTransformer的transform方法就会被调用，这里也就是com.sun.btrace.agent.Client#transformer了，该方法最后是调用了com.sun.btrace.agent.Client#instrument来完成真正的字节码修改工作，

 private byte[] instrument(Class clazz, String cname, byte[] target) {
            byte[] instrumentedCode;
            try {
                ClassWriter writer = InstrumentUtils.newClassWriter(target);
                ClassReader reader = new ClassReader(target);
                Instrumentor i = new Instrumentor(clazz, className,  btraceCode, onMethods, writer);
                InstrumentUtils.accept(reader, i);
                if (Main.isDebug() && !i.hasMatch()) {
                    Main.debugPrint("*WARNING* No method was matched for class " + cname); // NOI18N
                }
                instrumentedCode = writer.toByteArray();
            } catch (Throwable th) {
                Main.debugPrint(th);
                return null;
            }
            Main.dumpClass(className, cname, instrumentedCode);
            return instrumentedCode;
        }

上面使用的很多类的包名都是com.sun.btrace.org.objectweb.asm，BTrace使用了ASM来完成字节码的修改工作，具体细节暂时也不深究了。

最后简单总结一下，

1. BTrace脚本编译；
2. BTrace客户端使用Attach API attach到目标VM，并加载agent包；
3. agent打开socket来与客户端进行通信；
4. 客户端给agent发送InstrumentCommand，其中包含BTrace脚本编译后的字节码；
5. agent通过Attach API和ASM来完成满足BTrace脚本的字节码修改工作；