如何實現一個可用的javaagent

最近做了一個項目需要用javaagent方式對應用常用的組件（比如httpclient, 數據庫連接池等）進行調用追蹤和監控，並結合公司的分布式追蹤組件，將所有java應用的外部調用情況收集起來方便做系統分析和問題定位。項目定位和開源項目pinpoint比較像，但了解過pinpoint實現以后，發現其分布式追蹤和組件監控的邏輯耦合太過緊密，而且整個項目比較重，實現繁雜，不容易和公司的分布式追蹤組件結合起來，所以決定自己搞。這里暫且把這個項目取名叫dagent。

網上其實有很多文章介紹如何編寫javaagent，但往往介紹得非常簡單，只介紹premain的啟動機制，manifest如何編寫，但這類文章都沒有說明簡單實現的javaagent能否實際發揮作用，在實際的項目中可能會有哪些坑。所以，我想把這次項目過程中踩過的坑記錄下來，分享給需要的人。

ClassLoader之殤

首先得從spring boot的uber jar說起。所謂uber jar，就是一個all in one的可執行jar包。jar包中包含了Java應用運行所需要的代碼、資源以及依賴的jar包，直接執行java -jar xxx.jar即可啟動。對於web應用，spring boot還提供了嵌入式的web容器，無需部署tomcat服務器，應用部署運行特別方便。所以最近公司開始采用spring boot。

在測試dagent時發現，對於使用spring boot框架的應用，直接在IDE里面執行main方法運行dagent工作ok，一旦打成uber jar方式后加上dagent啟動，就會出現dagent中引用的第三方包中的類報ClassNotFoundException。

查看spring boot源碼，發現了原因所在。原來，由於uber jar將應用依賴的jar包以nested jar的方式打進包內，為了實現不解壓縮就啟動，spring boot使用自己的main類org.springframework.boot.loader.JarLauncher啟動應用，並自定義一個LaunchedURLClassLoader，再由它加載應用的main類。而LaunchedURLClassLoader在初始化classpath搜索路徑時特意把javaagent jar包排除在外，所以javaagent jar包中的類是不能被LaunchedURLClassLoader定義的，所以javaagent中的輔助類如果引用了某個第三方包中的類，而這個類是被LaunchedURLClassLoader定義的，簡單引用就會出現ClassNotFoundException。

所以，必須讓javaagent中用到的輔助類也由定義當前正在增強的Class的LaunchedURLClassLoader定義。

public interface ClassFileTransformer {
    byte[] transform(  
             ClassLoader         loader,
            String              className,
            Class<?>            classBeingRedefined,
            ProtectionDomain    protectionDomain,
            byte[]              classfileBuffer)
   throws IllegalClassFormatException;
}

更確切地說，應該讓每一個被ClassFileTransformer修改的類所用到的自定義以及第三方輔助類都由ClassFileTransformer#transform()方法第一個參數的ClassLoader定義。這樣，不管應用的ClassLoader采用了何種類查找策略，都可以保證輔助類可以正常加載到。

ClassLoader注入

那么，如何做到讓增強類所用到的輔助類都被同一個類定義呢。一種辦法是顯示地用對應的ClassLoader定義所有用到的輔助類，這樣需要手動注入所有輔助類，比較繁瑣；另一種辦法是將一組功能相關的輔助類打成jar包，注入到對應的ClassLoader中，這樣就不需要一個一個類手動注入。dagent采用的是第二種方案，將一組功能相關的增強輔助類做成一個插件，並打成一個jar包，然后在增強類的時候，將對應的jar包注入到當前執行增強的ClassLoader中。

不同類型的ClassLoader的注入方式有所不同，方法如下：

public class ClassInjector {
    private static Method DEFINE_CLASS;
    private static Method ADD_URL;

    static {
        try {
            DEFINE_CLASS = ClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class);
            DEFINE_CLASS.setAccessible(true);

            ADD_URL = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
            ADD_URL.setAccessible(true);
        } catch (NoSuchMethodException e) {
            throw new IllegalStateException(e);
        }
    }

    /**
     * 注入到非URLClassLoader的非引導類ClassLoader
     * @param classLoader
     * @param className
     * @param bytes 類定義
     * @throws InvocationTargetException
     * @throws IllegalAccessException
     */
    public static void defineClass(ClassLoader classLoader, String className, byte[] bytes) throws InvocationTargetException, IllegalAccessException {
        if (classLoader != null) {
            DEFINE_CLASS.invoke(classLoader, className, bytes, 0, bytes.length);
        }
    }

    /**
     * 注入到URLClassLoader類加載器
     * @param classLoader
     * @param url
     * @throws InvocationTargetException
     * @throws IllegalAccessException
     */
    public static void addURL(URLClassLoader classLoader, URL url) throws InvocationTargetException, IllegalAccessException {
        ADD_URL.invoke(classLoader, url);
    }

    /**
     * 注入到引導類加載器
     * @param instrumentation
     * @param jarFile
     */
    public static void addURL(Instrumentation instrumentation, JarFile jarFile) {
        instrumentation.appendToBootstrapClassLoaderSearch(jarFile);
    }
}

打包javaagent

采用上面的思路將javaagent用插件的方式實現，會導致每個插件都是一個jar包，不方便部署。可以用maven-assembly-plugin將javaagent核心代碼和插件打成一個jar包，agent加載時，再將jar包解壓，取出內嵌的插件包。解壓后的dagent包結構如下：

或者，也可以借鑒spring-boot uber jar的處理方式，自定義Jar包結構和URL handler，這樣就可以直接加載內嵌的jar包，不需要先解壓。