Repository: pinpoint-apm/pinpoint
Issue: RabbitMQ plugin may propagate incorrect async trace context to consumers
Labels: bug
module:plugin

#### Description
When **ConsumerDispatcher** receives a message, it delegates the message over to a consumer, which will handle the message in a different thread (**ExecutorService** in **ConsumerWorkService**).
In order to propagate trace over this thread boundary, **ConsumerDispatcher**'s interceptor injected async trace context directly to the consumer.
The consumer would then use the injected async trace context to continue tracing in a different thread.

Since the consumer is shared, a problem arises when the **ConsumerDispatcher** receives a new message before the consumer finishes handling the previous message (or the consumer might not have even started handling the previous message).
When this happens, the consumer will use the async trace context created from the new message, even if it's still handling the previous message.

Similar problem arises when the consumer's sampling rate is not 1 (100%). For example, if the sampling rate is set to 2 (50%), **ConsumerDispatcher** will create an async trace context for the first message and inject it into the consumer, and the consumer will use it to continue tracing. Then, when a second message arrives, **ConsumerDispatcher** will skip tracing, but the consumer still has the async trace context injected from the first message and will use it to continue tracing.

The resulting call stack for both cases would be jumbled up, with one of the transactions having both consumer traces, and the other having no consumer traces.

Having the **ConsumerDispatcher** inject async trace context into an object that is not shared would solve the problem.

Issue comments:

Pull request review comments:

Patch that resolved the issue:
diff --git a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/RabbitMQClientPlugin.java b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/RabbitMQClientPlugin.java
index d773d4fe46fd..e31fdbf35cfc 100644
--- a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/RabbitMQClientPlugin.java
+++ b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/RabbitMQClientPlugin.java
@@ -146,6 +146,15 @@ public byte[] doInTransform(Instrumentor instrumentor, ClassLoader loader, Strin
                 return target.toBytecode();
             }
         });
+        // Envelope - for asynchrnous trace propagation for consumers
+        transformTemplate.transform("com.rabbitmq.client.Envelope", new TransformCallback() {
+            @Override
+            public byte[] doInTransform(Instrumentor instrumentor, ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+                InstrumentClass target = instrumentor.getInstrumentClass(loader, className, classfileBuffer);
+                target.addField(AsyncContextAccessor.class.getName());
+                return target.toBytecode();
+            }
+        });
     }
 
     private void addAMQChannelEditor(final Filter<String> excludeExchangeFilter) {
@@ -187,7 +196,6 @@ private boolean addConsumerHandleDeliveryInterceptor(InstrumentClass target) thr
             return false;
         }
         handleDelivery.addScopedInterceptor("com.navercorp.pinpoint.plugin.rabbitmq.client.interceptor.ConsumerHandleDeliveryInterceptor", RabbitMQClientConstants.RABBITMQ_CONSUMER_SCOPE);
-        target.addField(AsyncContextAccessor.class.getName());
         return true;
     }
 
diff --git a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/ConsumerHandleDeliveryInterceptor.java b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/ConsumerHandleDeliveryInterceptor.java
index 9e154410a270..c02d68c0f1bc 100644
--- a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/ConsumerHandleDeliveryInterceptor.java
+++ b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/ConsumerHandleDeliveryInterceptor.java
@@ -1,29 +1,184 @@
 package com.navercorp.pinpoint.plugin.rabbitmq.client.interceptor;
 
+import com.navercorp.pinpoint.bootstrap.async.AsyncContextAccessor;
+import com.navercorp.pinpoint.bootstrap.async.AsyncContextAccessorUtils;
 import com.navercorp.pinpoint.bootstrap.context.*;
-import com.navercorp.pinpoint.bootstrap.interceptor.AsyncContextSpanEventSimpleAroundInterceptor;
+import com.navercorp.pinpoint.bootstrap.context.scope.TraceScope;
+import com.navercorp.pinpoint.bootstrap.interceptor.AroundInterceptor;
+import com.navercorp.pinpoint.bootstrap.logging.PLogger;
+import com.navercorp.pinpoint.bootstrap.logging.PLoggerFactory;
 import com.navercorp.pinpoint.plugin.rabbitmq.client.RabbitMQClientConstants;
+import com.rabbitmq.client.Consumer;
+import com.rabbitmq.client.Envelope;
 
 /**
  * @author Jinkai.Ma
  * @author Jiaqi Feng
  * @author HyunGil Jeong
  */
-public class ConsumerHandleDeliveryInterceptor extends AsyncContextSpanEventSimpleAroundInterceptor {
+public class ConsumerHandleDeliveryInterceptor implements AroundInterceptor {
+
+    private final PLogger logger = PLoggerFactory.getLogger(getClass());
+    private final boolean isDebug = logger.isDebugEnabled();
+    private static final String ASYNC_TRACE_SCOPE = AsyncContext.ASYNC_TRACE_SCOPE;
+
+    private final MethodDescriptor methodDescriptor;
 
     public ConsumerHandleDeliveryInterceptor(TraceContext traceContext, MethodDescriptor methodDescriptor) {
-        super(traceContext, methodDescriptor);
+        if (traceContext == null) {
+            throw new NullPointerException("traceContext must not be null");
+        }
+        if (methodDescriptor == null) {
+            throw new NullPointerException("methodDescriptor must not be null");
+        }
+        this.methodDescriptor = methodDescriptor;
     }
 
     @Override
-    protected void doInBeforeTrace(SpanEventRecorder recorder, AsyncContext asyncContext, Object target, Object[] args) {
+    public void before(Object target, Object[] args) {
+        if (!validate(target, args)) {
+            return;
+        }
+        final AsyncContext asyncContext = AsyncContextAccessorUtils.getAsyncContext(args[1]);
+        if (asyncContext == null) {
+            return;
+        }
+
+        final Trace trace = getAsyncTrace(asyncContext);
+        if (trace == null) {
+            return;
+        }
 
+        if (isDebug) {
+            logger.beforeInterceptor(target, args);
+        }
+
+        // entry scope.
+        entryAsyncTraceScope(trace);
+
+        try {
+            // trace event for default & async
+            final SpanEventRecorder recorder = trace.traceBlockBegin();
+            recorder.recordServiceType(RabbitMQClientConstants.RABBITMQ_CLIENT_INTERNAL);
+            recorder.recordApi(methodDescriptor);
+        } catch (Throwable t) {
+            if (logger.isWarnEnabled()) {
+                logger.warn("BEFORE error.", t);
+            }
+        }
     }
 
     @Override
-    protected void doInAfterTrace(SpanEventRecorder recorder, Object target, Object[] args, Object result, Throwable throwable) {
-        recorder.recordApi(methodDescriptor);
-        recorder.recordServiceType(RabbitMQClientConstants.RABBITMQ_CLIENT_INTERNAL);
-        recorder.recordException(throwable);
+    public void after(Object target, Object[] args, Object result, Throwable throwable) {
+        if (!validate(target, args)) {
+            return;
+        }
+        final AsyncContext asyncContext = AsyncContextAccessorUtils.getAsyncContext(args[1]);
+        if (asyncContext == null) {
+            return;
+        }
+
+        final Trace trace = asyncContext.currentAsyncTraceObject();
+        if (trace == null) {
+            return;
+        }
+
+        if (isDebug) {
+            logger.afterInterceptor(target, args, result, throwable);
+        }
+
+        // leave scope
+        if (!leaveAsyncTraceScope(trace)) {
+            if (logger.isWarnEnabled()) {
+                logger.warn("Failed to leave scope of async trace {}.", trace);
+            }
+            // delete unstable trace
+            deleteAsyncContext(trace, asyncContext);
+            return;
+        }
+
+        try {
+            final SpanEventRecorder recorder = trace.currentSpanEventRecorder();
+            if (throwable != null) {
+                recorder.recordException(throwable);
+            }
+        } catch (Throwable t) {
+            if (logger.isWarnEnabled()) {
+                logger.warn("AFTER error.", t);
+            }
+        } finally {
+            trace.traceBlockEnd();
+            if (isAsyncTraceDestination(trace)) {
+                deleteAsyncContext(trace, asyncContext);
+            }
+        }
+    }
+
+    private boolean validate(Object target, Object[] args) {
+        if (!(target instanceof Consumer)) {
+            return false;
+        }
+        if (args == null || args.length < 2) {
+            return false;
+        }
+        if (!(args[1] instanceof Envelope)) {
+            return false;
+        }
+        if (!(args[1] instanceof AsyncContextAccessor)) {
+            if (isDebug) {
+                logger.debug("Invalid args[1] object. Need accessor({}).", AsyncContextAccessor.class.getName());
+            }
+            return false;
+        }
+        return true;
+    }
+
+    private Trace getAsyncTrace(AsyncContext asyncContext) {
+        final Trace trace = asyncContext.continueAsyncTraceObject();
+        if (trace == null) {
+            if (logger.isWarnEnabled()) {
+                logger.warn("Failed to continue async trace. 'result is null'");
+            }
+            return null;
+        }
+        if (isDebug) {
+            logger.debug("getAsyncTrace() trace {}, asyncContext={}", trace, asyncContext);
+        }
+        return trace;
+    }
+
+    private void deleteAsyncContext(Trace trace, AsyncContext asyncContext) {
+        if (isDebug) {
+            logger.debug("Delete async trace {}.", trace);
+        }
+        trace.close();
+        asyncContext.close();
+    }
+
+    private void entryAsyncTraceScope(final Trace trace) {
+        final TraceScope scope = trace.getScope(ASYNC_TRACE_SCOPE);
+        if (scope != null) {
+            scope.tryEnter();
+        }
+    }
+
+    private boolean leaveAsyncTraceScope(final Trace trace) {
+        final TraceScope scope = trace.getScope(ASYNC_TRACE_SCOPE);
+        if (scope != null) {
+            if (scope.canLeave()) {
+                scope.leave();
+            } else {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private boolean isAsyncTraceDestination(final Trace trace) {
+        if (!trace.isAsync()) {
+            return false;
+        }
+        final TraceScope scope = trace.getScope(ASYNC_TRACE_SCOPE);
+        return scope != null && !scope.isActive();
     }
 }
diff --git a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/RabbitMQConsumerDispatchInterceptor.java b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/RabbitMQConsumerDispatchInterceptor.java
index 078c3763b665..6089e1d94757 100644
--- a/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/RabbitMQConsumerDispatchInterceptor.java
+++ b/plugins/rabbitmq/src/main/java/com/navercorp/pinpoint/plugin/rabbitmq/client/interceptor/RabbitMQConsumerDispatchInterceptor.java
@@ -92,10 +92,10 @@ public void before(Object target, Object[] args) {
             }
             SpanEventRecorder recorder = trace.traceBlockBegin();
             recorder.recordServiceType(RabbitMQClientConstants.RABBITMQ_CLIENT_INTERNAL);
-            // args[0] would be com.rabbitmq.client.Consumer, implementing AsyncContextAccessor via plugin
-            if (args[0] instanceof AsyncContextAccessor) {
+            // args[2] would be com.rabbitmq.client.Envelope, implementing AsyncContextAccessor via plugin
+            if (args[2] instanceof AsyncContextAccessor) {
                 AsyncContext asyncContext = recorder.recordNextAsyncContext();
-                ((AsyncContextAccessor) args[0])._$PINPOINT$_setAsyncContext(asyncContext);
+                ((AsyncContextAccessor) args[2])._$PINPOINT$_setAsyncContext(asyncContext);
             }
         } catch (Throwable th) {
             if (logger.isWarnEnabled()) {


Test patch:
diff --git a/agent/src/test/java/com/navercorp/pinpoint/plugin/jdk7/rabbitmq/RabbitMQTestRunner.java b/agent/src/test/java/com/navercorp/pinpoint/plugin/jdk7/rabbitmq/RabbitMQTestRunner.java
index 26216d166445..94b63e549607 100644
--- a/agent/src/test/java/com/navercorp/pinpoint/plugin/jdk7/rabbitmq/RabbitMQTestRunner.java
+++ b/agent/src/test/java/com/navercorp/pinpoint/plugin/jdk7/rabbitmq/RabbitMQTestRunner.java
@@ -36,16 +36,18 @@
 import org.junit.Assert;
 
 import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
 import java.util.concurrent.TimeUnit;
 
-
 /**
  * @author Jiaqi Feng
  * @author HyunGil Jeong
  */
 class RabbitMQTestRunner {
 
-    private final ConnectionFactory connectionFactory;
+    private static final Random RANDOM = new Random();
 
     RabbitMQTestRunner(ConnectionFactory connectionFactory) {
         if (connectionFactory == null) {
@@ -54,7 +56,14 @@ class RabbitMQTestRunner {
         this.connectionFactory = connectionFactory;
     }
 
+    private final ConnectionFactory connectionFactory;
+
     void runPushTest() throws Exception {
+        int numMessages = RANDOM.nextInt(10) + 1;
+        runPushTest(numMessages);
+    }
+
+    void runPushTest(int numMessages) throws Exception {
 
         final String message = "hello rabbit mq";
 
@@ -66,13 +75,15 @@ void runPushTest() throws Exception {
         producerChannel.queueDeclare(RabbitMQTestConstants.QUEUE_PUSH, false, false, false, null);
         producerChannel.queueBind(RabbitMQTestConstants.QUEUE_PUSH, RabbitMQTestConstants.EXCHANGE, RabbitMQTestConstants.ROUTING_KEY_PUSH);
 
-        AMQP.BasicProperties.Builder builder = new AMQP.BasicProperties.Builder();
-        producerChannel.basicPublish(RabbitMQTestConstants.EXCHANGE, RabbitMQTestConstants.ROUTING_KEY_PUSH, false, false, builder.appId("test").build(), message.getBytes());
+        for (int i = 0; i < numMessages; i++) {
+            AMQP.BasicProperties.Builder builder = new AMQP.BasicProperties.Builder();
+            producerChannel.basicPublish(RabbitMQTestConstants.EXCHANGE, RabbitMQTestConstants.ROUTING_KEY_PUSH, false, false, builder.appId("test").build(), message.getBytes());
+        }
 
         producerChannel.close();
         producerConnection.close();
 
-        //comsumer side
+        // consumer side
         final Connection consumerConnection = connectionFactory.newConnection();
         final Channel consumerChannel = consumerConnection.createChannel();
         final String remoteAddress = consumerConnection.getAddress().getHostAddress() + ":" + consumerConnection.getPort();
@@ -82,17 +93,26 @@ void runPushTest() throws Exception {
         TestConsumer<String> consumer = new TestConsumer<String>(consumerChannel, MessageConverter.FOR_TEST);
         consumerChannel.basicConsume(RabbitMQTestConstants.QUEUE_PUSH, true, consumer);
 
-        // wait consumer
-        Assert.assertEquals(message, consumer.getMessage(10, TimeUnit.SECONDS));
+        List<String> actualMessages = new ArrayList<String>(numMessages);
+        for (int i = 0; i < numMessages; i++) {
+            actualMessages.add(consumer.getMessage(10, TimeUnit.SECONDS));
+        }
+
+        Assert.assertEquals(numMessages, actualMessages.size());
+        for (String actualMessage : actualMessages) {
+            Assert.assertEquals(message, actualMessage);
+        }
 
         consumerChannel.close();
         consumerConnection.close();
 
         PluginTestVerifier verifier = PluginTestVerifierHolder.getInstance();
         // Wait till all traces are recorded (consumer traces are recorded from another thread)
-        awaitAndVerifyTraceCount(verifier, 6, 5000L);
+        int expectedTraceCountPerMessage = 6;
+        awaitAndVerifyTraceCount(verifier, expectedTraceCountPerMessage * numMessages, 5000L);
 
         verifier.printCache();
+
         Class<?> producerChannelClass = producerChannel.getClass();
         Method channelBasicPublish = producerChannelClass.getDeclaredMethod("basicPublish", String.class, String.class, boolean.class, boolean.class, AMQP.BasicProperties.class, byte[].class);
         ExpectedTrace channelBasicPublishTrace = Expectations.event(
@@ -128,13 +148,17 @@ void runPushTest() throws Exception {
         ExpectedTrace markTrace = Expectations.event(
                 ServiceType.INTERNAL_METHOD.getName(),
                 propagationMarkerMark);
-        verifier.verifyTrace(
-                channelBasicPublishTrace,
-                rabbitMqConsumerInvocationTrace,
-                consumerDispatcherHandleDeliveryTrace,
-                asynchronousInvocationTrace,
-                consumerHandleDeliveryTrace,
-                markTrace);
+
+        for (int i = 0; i < numMessages; i++) {
+            verifier.verifyDiscreteTrace(channelBasicPublishTrace);
+            verifier.verifyDiscreteTrace(
+                    rabbitMqConsumerInvocationTrace,
+                    Expectations.async(
+                            consumerDispatcherHandleDeliveryTrace,
+                            asynchronousInvocationTrace,
+                            consumerHandleDeliveryTrace,
+                            markTrace));
+        }
         verifier.verifyTraceCount(0);
     }