From 32c64c46c343555a111c52040644d544a3499a30 Mon Sep 17 00:00:00 2001
From: Andrea Marziali <andrea.marziali@datadoghq.com>
Date: Fri, 6 Dec 2024 10:19:09 +0100
Subject: [PATCH] Protect mdb from instrumenting multiple time the same event
 (#8062)

---
 .../jakarta/jms/MDBMessageConsumerInstrumentation.java      | 6 ++++++
 .../jms/MDBMessageConsumerInstrumentation.java              | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/dd-java-agent/instrumentation/jakarta-jms/src/main/java/datadog/trace/instrumentation/jakarta/jms/MDBMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jakarta-jms/src/main/java/datadog/trace/instrumentation/jakarta/jms/MDBMessageConsumerInstrumentation.java
index 1217e383aad..b9b651110bf 100644
--- a/dd-java-agent/instrumentation/jakarta-jms/src/main/java/datadog/trace/instrumentation/jakarta/jms/MDBMessageConsumerInstrumentation.java
+++ b/dd-java-agent/instrumentation/jakarta-jms/src/main/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/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java
index 2c55e314141..c5f2233cabb 100644
--- a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java
+++ b/dd-java-agent/instrumentation/jms/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();