如何写一个Skywalking trace插件

javaagent 原理

美团技术团队-Java 动态调试技术原理及实践

类图

实现

ConsumeMessageConcurrentlyInstrumentation

public class ConsumeMessageConcurrentlyInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
    // 需要增强的类
    private static final String ENHANCE_CLASS = "com.aliyun.openservices.shade.com.alibaba.rocketmq.client.consumer.listener.MessageListenerConcurrently";
    // 需要增强的方法
    private static final String CONSUMER_MESSAGE_METHOD = "consumeMessage";
    // 增加的方法对应的拦截器
    private static final String INTERCEPTOR_CLASS = "org.apache.skywalking.apm.plugin.ons.v1.MessageConcurrentlyConsumeInterceptor";

    // 构造器不需要拦截
    @Override
    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
        return new ConstructorInterceptPoint[0];
    }

    @Override
    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
        return new InstanceMethodsInterceptPoint[] {
            // 新增一个拦截器
            new InstanceMethodsInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getMethodsMatcher() {
                    // 方法匹配
                    return named(CONSUMER_MESSAGE_METHOD);
                }

                @Override
                public String getMethodsInterceptor() {
                    return INTERCEPTOR_CLASS;
                }

                @Override
                public boolean isOverrideArgs() {
                    return false;
                }
            }
        };
    }

    @Override
    protected ClassMatch enhanceClass() {
        // 需要增强的类
        return HierarchyMatch.byHierarchyMatch(new String[] {ENHANCE_CLASS});
    }
}

AbstractMessageConsumeInterceptor

public abstract class AbstractMessageConsumeInterceptor implements InstanceMethodsAroundInterceptor {

    public static final String CONSUMER_OPERATION_NAME_PREFIX = "OnsRocketMQ/";

    // 在方法前增强
    @Override
    public final void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
        // 拿到方法参数，转换成消息列表
        List<MessageExt> msgs = (List<MessageExt>) allArguments[0];

        // 从消息中中获取TraceId等Context信息
        ContextCarrier contextCarrier = getContextCarrierFromMessage(msgs.get(0));
        
        // 创建一个entry span
        AbstractSpan span = ContextManager.createEntrySpan(CONSUMER_OPERATION_NAME_PREFIX + msgs.get(0)
                                                                                                .getTopic() + "/Consumer", contextCarrier);
        span.setComponent(ComponentsDefine.ROCKET_MQ_CONSUMER);
        SpanLayer.asMQ(span);
        for (int i = 1; i < msgs.size(); i++) {
            ContextManager.extract(getContextCarrierFromMessage(msgs.get(i)));
        }

    }

    // 异常处理
    @Override
    public final void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, Throwable t) {
        ContextManager.activeSpan().log(t);
    }

    private ContextCarrier getContextCarrierFromMessage(MessageExt message) {
        ContextCarrier contextCarrier = new ContextCarrier();

        CarrierItem next = contextCarrier.items();
        while (next.hasNext()) {
            next = next.next();
            next.setHeadValue(message.getUserProperty(next.getHeadKey()));
        }

        return contextCarrier;
    }
}

MessageConcurrentlyConsumeInterceptor

public class MessageConcurrentlyConsumeInterceptor extends AbstractMessageConsumeInterceptor {
    // 在方法后处理
    @Override
    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
        Object ret) throws Throwable {
        // 获取消费状态
        ConsumeConcurrentlyStatus status = (ConsumeConcurrentlyStatus) ret;
        if (status == ConsumeConcurrentlyStatus.RECONSUME_LATER) {
			// 消费状态为重试，则设置span出现错误
            AbstractSpan activeSpan = ContextManager.activeSpan();
            activeSpan.errorOccurred();
            Tags.MQ_STATUS.set(activeSpan, status.name());
        }
        // 停止span
        ContextManager.stopSpan();
        return ret;
    }
}

项目：apm-ons-1.x-plugin

参考文档

1. apm-ons-1.x-plugin
2. 美团技术团队-Java 动态调试技术原理及实践

分享并记录所学所见