Protect mdb from instrumenting multiple time the same event (#8062)

DataDog · Dec 6, 2024 · 32c64c4 · 32c64c4
1 parent 2b24697
commit 32c64c4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/...ain/java/datadog/trace/instrumentation/jakarta/jms/MDBMessageConsumerInstrumentation.java b/...ain/java/datadog/trace/instrumentation/jakarta/jms/MDBMessageConsumerInstrumentation.java
@@ -19,11 +19,13 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import jakarta.jms.Destination;
 import jakarta.jms.JMSException;
 import jakarta.jms.Message;
+import jakarta.jms.MessageListener;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -72,6 +74,9 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class MDBAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope methodEnter(@Advice.Argument(0) final Message message) {
+      if (CallDepthThreadLocalMap.incrementCallDepth(MessageListener.class) > 0) {
+        return null;
+      }
       AgentSpan.Context propagatedContext = propagate().extract(message, GETTER);
       AgentSpan span = startSpan(JMS_CONSUME, propagatedContext);
       CONSUMER_DECORATE.afterStart(span);
@@ -93,6 +98,7 @@ public static AgentScope methodEnter(@Advice.Argument(0) final Message message)
     public static void methodExit(
         @Advice.Enter AgentScope scope, @Advice.Thrown final Throwable throwable) {
       if (null != scope) {
+        CallDepthThreadLocalMap.reset(MessageListener.class);
         CONSUMER_DECORATE.onError(scope, throwable);
         scope.close();
         scope.span().finish();

diff --git a/...ms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java b/...ms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java
@@ -19,11 +19,13 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import javax.jms.Destination;
 import javax.jms.JMSException;
 import javax.jms.Message;
+import javax.jms.MessageListener;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -72,6 +74,9 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class MDBAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope methodEnter(@Advice.Argument(0) final Message message) {
+      if (CallDepthThreadLocalMap.incrementCallDepth(MessageListener.class) > 0) {
+        return null;
+      }
       AgentSpan.Context propagatedContext = propagate().extract(message, GETTER);
       AgentSpan span = startSpan(JMS_CONSUME, propagatedContext);
       CONSUMER_DECORATE.afterStart(span);
@@ -93,6 +98,7 @@ public static AgentScope methodEnter(@Advice.Argument(0) final Message message)
     public static void methodExit(
         @Advice.Enter AgentScope scope, @Advice.Thrown final Throwable throwable) {
       if (null != scope) {
+        CallDepthThreadLocalMap.reset(MessageListener.class);
         CONSUMER_DECORATE.onError(scope, throwable);
         scope.close();
         scope.span().finish();