《arthas源码分析》字节码增强（一）

前言

我们都知道aop的实现原理，是在方法的前后做了一些事情，arthas的一些命令的实现原理也是如此，比如：watch 、trace 等命令。那本质来说，是通过字节码增强技术实现的。

本篇将对字节码的增强过程进行详细分析。

阅读本篇需要对java agent，以及asm有一定的了解。

源码分析

EnhancerCommand#process

每个命令的入口都是process()方法

@Override
public void process(final CommandProcess process) {
    // ctrl-C support 
    process.interruptHandler(new CommandInterruptHandler(process));
    // q exit support 
    process.stdinHandler(new QExitHandler(process));
    // start to enhance 
    enhance(process);
}

EnhancerCommand#enhance

对目标类进行增强

protected void enhance(CommandProcess process) {
    Session session = process.session();
    //如果已加锁，直接返回
    if (!session.tryLock()) {
        process.write("someone else is enhancing classes, pls. wait.\n");
        process.end();
        return;
    }
    //获取锁
    int lock = session.getLock();
    try {
        //从session中获取Instrumentation
        Instrumentation inst = session.getInstrumentation();
        //获取当前命令的通知监听器，比如：WatchAdviceListener
        AdviceListener listener = getAdviceListener(process);
        if (listener == null) {
            warn(process, "advice listener is null");
            return;
        }
        
        //是否跳过JDK，默认是跳过，如果是trace命令的话，根据其参数决定是否跳过
        boolean skipJDKTrace = false;
        if(listener instanceof AbstractTraceAdviceListener) {
            skipJDKTrace = ((AbstractTraceAdviceListener) listener).getCommand().isSkipJDKTrace();
        }

        //对目标类进行增强
        EnhancerAffect effect = Enhancer.enhance(inst, lock, listener instanceof InvokeTraceable,
                skipJDKTrace, getClassNameMatcher(), getMethodNameMatcher());

        //增强的类或者方法的个数为0，则直接返回
        if (effect.cCnt() == 0 || effect.mCnt() == 0) {
            // no class effected
            // might be method code too large
            //响应信息
            process.write("No class or method is affected, try:\n"
                          + "1. sm CLASS_NAME METHOD_NAME to make sure the method you are tracing actually exists (it might be in your parent class).\n"
                          + "2. reset CLASS_NAME and try again, your method body might be too large.\n"
                          + "3. check arthas log: " + LogUtil.LOGGER_FILE + "\n"
                          + "4. visit https://github.com/alibaba/arthas/issues/47 for more details.\n");
            //结束命令
            process.end();
            return;
        }

        // 这里做个补偿,如果在enhance期间,unLock被调用了,则补偿性放弃
        if (session.getLock() == lock) {
            // 注册监听器
            process.register(lock, listener);
            // 如果命令在前台运行，则响应提示信息："Press Q or Ctrl+C to abort."
            if (process.isForeground()) {
                process.echoTips(Constants.Q_OR_CTRL_C_ABORT_MSG + "\n");
            }
        }

        //响应增强影响
        process.write(effect + "\n");
    } catch (UnmodifiableClassException e) {
        logger.error(null, "error happens when enhancing class", e);
    } finally {
        if (session.getLock() == lock) {
            // enhance结束后解锁
            process.session().unLock();
        }
    }
}

Enhancer#enhance

对目标类进行增强，核心是通过如下Api对目标类进行增强
Instrumentation#addTransformer ：添加字节码文件转换器
Instrumentation#retransformClasses ：转换字节码文件，对目标类进行增强。

说明：
1、Enhancer 的构造方法均被private修饰，所以其只能通过其enhance()方法进行创建。
2、Enhancer实现了ClassFileTransformer，其transform()方法定义了增强的核心逻辑。
3、增强的两种方式（批量增强和单个增强）
两者所使用的API是一致的，不同点在于，单个增强，可以打印更详细的日志信息。

public static synchronized EnhancerAffect enhance(
        final Instrumentation inst,
        final int adviceId,
        final boolean isTracing,
        final boolean skipJDKTrace,
        final Matcher classNameMatcher,
        final Matcher methodNameMatcher) throws UnmodifiableClassException {

    //增强影响
    final EnhancerAffect affect = new EnhancerAffect();

    //查找目标类。
    final Set<Class<?>> enhanceClassSet = GlobalOptions.isDisableSubClass
            ? SearchUtils.searchClass(inst, classNameMatcher)
            : SearchUtils.searchSubClass(inst, SearchUtils.searchClass(inst, classNameMatcher));

    // 过滤掉无法被增强的类
    filter(enhanceClassSet);

    //创建增强器
    final Enhancer enhancer = new Enhancer(adviceId, isTracing, skipJDKTrace, enhanceClassSet, methodNameMatcher, affect);
    try {
        //添加.class文件转换器: ClassFileTransformer
        inst.addTransformer(enhancer, true);

        //批量增强
        if (GlobalOptions.isBatchReTransform) {
            final int size = enhanceClassSet.size();
            final Class<?>[] classArray = new Class<?>[size];
            arraycopy(enhanceClassSet.toArray(), 0, classArray, 0, size);
            if (classArray.length > 0) {
                //增强目标类
                inst.retransformClasses(classArray);
                logger.info("Success to batch transform classes: " + Arrays.toString(classArray));
            }
        }
        //一个一个地增强
        else {
            for (Class<?> clazz : enhanceClassSet) {
                try {
                    inst.retransformClasses(clazz);
                    logger.info("Success to transform class: " + clazz);
                } catch (Throwable t) {
                    logger.warn("retransform {} failed.", clazz, t);
                    if (t instanceof UnmodifiableClassException) {
                        throw (UnmodifiableClassException) t;
                    } else if (t instanceof RuntimeException) {
                        throw (RuntimeException) t;
                    } else {
                        throw new RuntimeException(t);
                    }
                }
            }
        }
    } finally {
        //增强完毕，移除目标类的增强器，无需多次，也不可以多次增强
        inst.removeTransformer(enhancer);
    }

    return affect;
}

Enhancer#transform

目标类的核心增强逻辑。

说明：
1、这里使用到了ASM框架，我也会在接下来的篇幅中有所介绍
2、通过重写ClassVisitor、MethodVisitor相关方法，达到增强目标类的目的，详情参见：AdviceWeaver。

@Override
public byte[] transform(final ClassLoader inClassLoader, String className, Class<?> classBeingRedefined,
                ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
    try {
        // 这里要再次过滤一次，为啥？因为在transform的过程中，有可能还会再诞生新的类
        // 所以需要将之前需要转换的类集合传递下来，再次进行判断
        if (!matchingClasses.contains(classBeingRedefined)) {
            return null;
        }

        final ClassReader cr;

        // 首先先检查是否在缓存中存在Class字节码
        // 因为要支持多人协作,存在多人同时增强的情况
        final byte[] byteOfClassInCache = classBytesCache.get(classBeingRedefined);
        if (null != byteOfClassInCache) {
            cr = new ClassReader(byteOfClassInCache);
        }

        // 如果没有命中缓存,则从原始字节码开始增强
        else {
            cr = new ClassReader(classfileBuffer);
        }

        // 字节码增强
        final ClassWriter cw = new ClassWriter(cr, COMPUTE_FRAMES | COMPUTE_MAXS) {

            /*
             * 注意，为了自动计算帧的大小，有时必须计算两个类共同的父类。
             * 缺省情况下，ClassWriter将会在getCommonSuperClass方法中计算这些，通过在加载这两个类进入虚拟机时，使用反射API来计算。
             * 但是，如果你将要生成的几个类相互之间引用，这将会带来问题，因为引用的类可能还不存在。
             * 在这种情况下，你可以重写getCommonSuperClass方法来解决这个问题。
             *
             * 通过重写 getCommonSuperClass() 方法，更正获取ClassLoader的方式，改成使用指定ClassLoader的方式进行。
             * 规避了原有代码采用Object.class.getClassLoader()的方式
             */
            @Override
            protected String getCommonSuperClass(String type1, String type2) {
                Class<?> c, d;
                final ClassLoader classLoader = inClassLoader;
                try {
                    c = Class.forName(type1.replace('/', '.'), false, classLoader);
                    d = Class.forName(type2.replace('/', '.'), false, classLoader);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
                if (c.isAssignableFrom(d)) {
                    return type1;
                }
                if (d.isAssignableFrom(c)) {
                    return type2;
                }
                if (c.isInterface() || d.isInterface()) {
                    return "java/lang/Object";
                } else {
                    do {
                        c = c.getSuperclass();
                    } while (!c.isAssignableFrom(d));
                    return c.getName().replace('.', '/');
                }
            }

        };

        // 生成增强字节码
        // AdviceWeaver：通过重写ASM相关方法，真正实现增强逻辑。
        cr.accept(new AdviceWeaver(adviceId, isTracing, skipJDKTrace, cr.getClassName(), methodNameMatcher, affect,
                        cw), EXPAND_FRAMES);
        final byte[] enhanceClassByteArray = cw.toByteArray();

        //将字节码推入缓存
        classBytesCache.put(classBeingRedefined, enhanceClassByteArray);

        // dump the class
        //如果需要dump，则将增强后的字节码刷写到磁盘
        dumpClassIfNecessary(className, enhanceClassByteArray, affect);

        // 成功计数(记录增强的类的个数)
        affect.cCnt(1);

        // 派遣间谍
        try {
            spy(inClassLoader);
        } catch (Throwable t) {
            logger.warn("print spy failed. classname={};loader={};", className, inClassLoader, t);
            throw t;
        }

        return enhanceClassByteArray;
    } catch (Throwable t) {
        logger.warn("transform loader[{}]:class[{}] failed.", inClassLoader, className, t);
    }

    return null;
}

Enhancer#dumpClassIfNecessary

dump增强类的字节码

 private static void dumpClassIfNecessary(String className, byte[] data, EnhancerAffect affect) {
        //如果不需要dump直接返回
        if (!GlobalOptions.isDump) {
            return;
        }
        //字节码文件
        final File dumpClassFile = new File("./arthas-class-dump/" + className + ".class");
        //.class文件目录
        final File classPath = new File(dumpClassFile.getParent());

        //如果不存在，则创建目录
        if (!classPath.mkdirs()
                && !classPath.exists()) {
            logger.warn("create dump classpath:{} failed.", classPath);
            return;
        }

        // 将类字节码写入磁盘
        try {
            FileUtils.writeByteArrayToFile(dumpClassFile, data);
            affect.getClassDumpFiles().add(dumpClassFile);
        } catch (IOException e) {
            logger.warn("dump class:{} to file {} failed.", className, dumpClassFile, e);
        }
    }

派遣间谍

1、Spy类的作用：主要是记录advice的钩子的引用，在增强的时候，会通过反射获取。
2、spy()方法的作用：在arthas启动的时候，将各种钩子设置到了自定义类加载器ArthasClassLoader，这里将其取出，设置到目标类的classloader当中。

Enhancer#spy

private void spy(final ClassLoader targetClassLoader) throws Exception {
    if (targetClassLoader == null) {
        // 增强JDK自带的类,targetClassLoader为null
        return;
    }
    // 因为 Spy 是被bootstrap classloader加载的，所以一定可以被找到，如果找不到的话，说明应用方的classloader实现有问题
    Class<?> spyClass = targetClassLoader.loadClass(Constants.SPY_CLASSNAME);

    final ClassLoader arthasClassLoader = Enhancer.class.getClassLoader();

    // 初始化间谍, AgentLauncher会把各种hook设置到ArthasClassLoader当中
    // 这里我们需要把这些hook取出来设置到目标classloader当中
    Method initMethod = spyClass.getMethod("init", ClassLoader.class, Method.class,
            Method.class, Method.class, Method.class, Method.class, Method.class);
    initMethod.invoke(null, arthasClassLoader,
            FieldUtils.getField(spyClass, "ON_BEFORE_METHOD").get(null),
            FieldUtils.getField(spyClass, "ON_RETURN_METHOD").get(null),
            FieldUtils.getField(spyClass, "ON_THROWS_METHOD").get(null),
            FieldUtils.getField(spyClass, "BEFORE_INVOKING_METHOD").get(null),
            FieldUtils.getField(spyClass, "AFTER_INVOKING_METHOD").get(null),
            FieldUtils.getField(spyClass, "THROW_INVOKING_METHOD").get(null));
}

SPY

public class Spy {
    // -- 各种Advice的钩子引用 --
    public static volatile Method ON_BEFORE_METHOD;
    public static volatile Method ON_RETURN_METHOD;
    public static volatile Method ON_THROWS_METHOD;
    public static volatile Method BEFORE_INVOKING_METHOD;
    public static volatile Method AFTER_INVOKING_METHOD;
    public static volatile Method THROW_INVOKING_METHOD;

    /**
     * arthas's classloader 引用
     */
    public static volatile ClassLoader CLASSLOADER;

    /**
     * 代理重设方法
     */
    public static volatile Method AGENT_RESET_METHOD;

    /**
     * 用于普通的间谍初始化
     */
    public static void init(
            ClassLoader classLoader,
            Method onBeforeMethod,
            Method onReturnMethod,
            Method onThrowsMethod,
            Method beforeInvokingMethod,
            Method afterInvokingMethod,
            Method throwInvokingMethod) {
        CLASSLOADER = classLoader;
        ON_BEFORE_METHOD = onBeforeMethod;
        ON_RETURN_METHOD = onReturnMethod;
        ON_THROWS_METHOD = onThrowsMethod;
        BEFORE_INVOKING_METHOD = beforeInvokingMethod;
        AFTER_INVOKING_METHOD = afterInvokingMethod;
        THROW_INVOKING_METHOD = throwInvokingMethod;
    }

    /**
     * 用于启动线程初始化
     */
    public static void initForAgentLauncher(
            ClassLoader classLoader,
            Method onBeforeMethod,
            Method onReturnMethod,
            Method onThrowsMethod,
            Method beforeInvokingMethod,
            Method afterInvokingMethod,
            Method throwInvokingMethod,
            Method agentResetMethod) {
        CLASSLOADER = classLoader;
        ON_BEFORE_METHOD = onBeforeMethod;
        ON_RETURN_METHOD = onReturnMethod;
        ON_THROWS_METHOD = onThrowsMethod;
        BEFORE_INVOKING_METHOD = beforeInvokingMethod;
        AFTER_INVOKING_METHOD = afterInvokingMethod;
        THROW_INVOKING_METHOD = throwInvokingMethod;
        AGENT_RESET_METHOD = agentResetMethod;
    }

    /**
     * Clean up the reference to com.taobao.arthas.agent.AgentLauncher$1
     * to avoid classloader leak.
     */
    public static void destroy() {
        CLASSLOADER = null;
        ON_BEFORE_METHOD = null;
        ON_RETURN_METHOD = null;
        ON_THROWS_METHOD = null;
        BEFORE_INVOKING_METHOD = null;
        AFTER_INVOKING_METHOD = null;
        THROW_INVOKING_METHOD = null;
        // clear the reference to ArthasClassLoader in AgentLauncher
        if (AGENT_RESET_METHOD != null) {
            try {
                AGENT_RESET_METHOD.invoke(null);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        AGENT_RESET_METHOD = null;
    }

}