Skywalking中RabbitMQ消费链路被隔断

skywalking官方是支持rabbitMQ探针的，理想状态下是mq消费的链路是被包在一个端点下的，
但实际情况是消费过程中如若调用其他链路会被隔断为多个链路，如下：

标红部分应该被放在一个链路中。继而研究源码，此处省略n多熟悉skywalking的过程…
在apm-sniffer -> apm-sdk-plugin -> rabbitmq5.x-plugin 中

public class RabbitMQConsumerInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
    public static final String INTERCEPTOR_CLASS = "org.apache.skywalking.apm.plugin.rabbitmq.RabbitMQConsumerInterceptor";
    public static final String ENHANCE_CLASS_PRODUCER = "com.rabbitmq.client.impl.ConsumerDispatcher";
    public static final String ENHANCE_METHOD_DISPATCH = "handleDelivery";
    public static final String INTERCEPTOR_CONSTRUCTOR = "org.apache.skywalking.apm.plugin.rabbitmq.RabbitMQProducerAndConsumerConstructorInterceptor";
    //........
 }

RabbitMQConsumerInstrumentation这个类声明了agent探针对具体的某个类的某个方法进行增强，继续往下

final class ConsumerDispatcher {
	// .....
    public void handleDelivery(final Consumer delegate,
                           final String consumerTag,
                           final Envelope envelope,
                           final AMQP.BasicProperties properties,
                           final byte[] body) throws IOException {
    executeUnlessShuttingDown(
        new Runnable() {
        @Override
        public void run() {
            try {
                delegate.handleDelivery(consumerTag,
                        envelope,
                        properties,
                        body);
            } catch (Throwable ex) {
                connection.getExceptionHandler().handleConsumerException(
                        channel,
                        ex,
                        delegate,
                        consumerTag,
                        "handleDelivery");
            }
        }
    });
    }
    private void executeUnlessShuttingDown(Runnable r) {
    if (!this.shuttingDown) execute(r);
    }
 
    private void execute(Runnable r) {
       checkShutdown();
       this.workService.addWork(this.channel, r);
    }
    //......
}

handleDelivery方法是rabbitMq中client端监听到消息后调用的方法，源码可见，直接生成了run方法然后丢到线程池中。
问题就来了，agent收集trace信息依赖于threadlocal变量将一个线程内调用的方法串起来组成一条trace，而delegate.handleDelivery()方法是具体的消费逻辑，并不会被执行ConsumerDispatcher.handleDelivery()的线程去执行，所以会导致trace被隔断。

现在已经找到问题所在了，解决思路如下：

1. 想办法将ConsumerDispatcher.handleDelivery()的线程变量带到run() 中
2. 修改agent增强点，不增强监听的方法，转而去增强消费的方法

第1种思路可以使用agent对run方法修改，但agent只能增强无法替换源方法，只能修改rabbitmq的源码，将线程变量带入，但过于复杂，无法实现。

继而研究第2种思路。

public interface MessageListener {
   void onMessage(Message message);
}

MessageListener.onMessage()是spring与rabbitmq整合后提供的消费接口，供用户自己去实现具体消费逻辑，但agent无法对interface或者abstract方法进行增强，

所以只能转而求其次，寻找实体方法,对messageListener打断点进行溯源：

在org.springframework.amqp.rabbit.listener.AbstractMessageListenerContainer抽象类中定义了执行消费的逻辑，我们可以在此处下手，继续往下




好了，就是这里了，org.springframework.amqp.rabbit.listener.AbstractMessageListenerContainer#invokeListener(Channel channel, Message message)
这里的listener是接口，spring会根据不同的message调用不同的listener的onMessage()方法，
既然已经找到了切点，就可以开始修改rabbitmq-plugin的源码了

public class RabbitMQConsumerInvokeInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {

    public static final String INTERCEPTOR_CLASS = "org.apache.skywalking.apm.plugin.rabbitmq.RabbitMQConsumerInvokeInterceptor";
    public static final String ENHANCE_CLASS_PRODUCER = "org.springframework.amqp.rabbit.listener.AbstractMessageListenerContainer";
    public static final String ENHANCE_METHOD_DISPATCH = "invokeListener";
    public static final String INTERCEPTOR_CONSTRUCTOR = "org.apache.skywalking.apm.plugin.rabbitmq.RabbitMQProducerAndConsumerConstructorInterceptor";
    @Override
    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
        return new ConstructorInterceptPoint[] {
            new ConstructorInterceptPoint() {
                    @Override public ElementMatcher<MethodDescription> getConstructorMatcher() {
                        return takesArgumentWithType(0,"com.rabbitmq.client.impl.AMQConnection");
                    }

                    @Override public String getConstructorInterceptor() {
                        return INTERCEPTOR_CONSTRUCTOR;
                    }
                }
        };
    }

    @Override
    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
        return new InstanceMethodsInterceptPoint[] {
            new InstanceMethodsInterceptPoint() {
                    @Override public ElementMatcher<MethodDescription> getMethodsMatcher() {
                        return named(ENHANCE_METHOD_DISPATCH).and(takesArgumentWithType(1,"org.springframework.amqp.core.Message"));
                    }


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

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

    @Override
    protected ClassMatch enhanceClass() {
        return MultiClassNameMatch.byMultiClassMatch(ENHANCE_CLASS_PRODUCER);
    }
}

修改拦截方法

public class RabbitMQConsumerInvokeInterceptor implements InstanceMethodsAroundInterceptor {
    public static final String OPERATE_NAME_PREFIX = "RabbitMQ/";
    public static final String CONSUMER_OPERATE_NAME_SUFFIX = "/Consumer/invoke/";
 
    @Override
    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
        ContextCarrier contextCarrier = new ContextCarrier();

        Connection connection = ((Channel)allArguments[0]).getConnection();
        String url =  connection.getAddress().toString().replace("/","") + ":" + connection.getPort();
        Message message = (Message) allArguments[1];
        MessageProperties msgProperties = message.getMessageProperties();
        String msgTopic = msgProperties.getReceivedRoutingKey();
        String msgContent = new String(message.getBody());
        AbstractSpan activeSpan = ContextManager.createEntrySpan(OPERATE_NAME_PREFIX + "Topic/" + msgTopic +
                CONSUMER_OPERATE_NAME_SUFFIX + msgContent, null).start(System.currentTimeMillis());

        Tags.MQ_BROKER.set(activeSpan, url);
        Tags.MQ_TOPIC.set(activeSpan, msgProperties.getReceivedExchange());
        Tags.MQ_QUEUE.set(activeSpan, msgProperties.getReceivedRoutingKey());
        activeSpan.setComponent(ComponentsDefine.RABBITMQ_CONSUMER);
        SpanLayer.asMQ(activeSpan);
        CarrierItem next = contextCarrier.items();
        while (next.hasNext()) {
            next = next.next();
            if (msgProperties.getHeaders() != null && msgProperties.getHeader(next.getHeadKey()) != null) {
                next.setHeadValue(msgProperties.getHeader(next.getHeadKey()).toString());
            }
        }
        ContextManager.extract(contextCarrier);
    }
}

重新编译、打包，最终效果如下：

注：此方法基于skywalking官方的rabbitmq-5.x-plugin插件修改
附git连接：https://github.com/To404/skywalking-rabbitmq-plugin