Add recursive loop protection

DataDog · Oct 7, 2024 · afe965a · afe965a
1 parent a4fe058
commit afe965a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumReaderInstrumentation.java b/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumReaderInstrumentation.java
@@ -8,6 +8,7 @@
 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.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -48,11 +49,20 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   }
 
   public static class GenericDatumReaderAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void trackDepth() {
+      CallDepthThreadLocalMap.incrementCallDepth(GenericDatumReader.class);
+    }
+
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(@Advice.This GenericDatumReader reader) {
       if (!Config.get().isDataStreamsEnabled()) {
         return;
       }
+      final int callDepth = CallDepthThreadLocalMap.decrementCallDepth(GenericDatumReader.class);
+      if (callDepth > 0) {
+        return;
+      }
       AgentSpan span = activeSpan();
       if (span == null) {
         return;

diff --git a/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumWriterInstrumentation.java b/...o/src/main/java/datadog/trace/instrumentation/avro/GenericDatumWriterInstrumentation.java
@@ -9,12 +9,14 @@
 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.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericDatumWriter;
 
 @AutoService(InstrumenterModule.class)
 public final class GenericDatumWriterInstrumentation extends InstrumenterModule.Tracing
@@ -49,11 +51,23 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   }
 
   public static class GenericDatumWriterAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.This GenericDatumWriter reader) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(GenericDatumWriter.class);
+      if (callDepth > 0) {
+        return;
+      }
+    }
+
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(@Advice.FieldValue("root") Schema root) {
       if (!Config.get().isDataStreamsEnabled()) {
         return;
       }
+      final int callDepth = CallDepthThreadLocalMap.decrementCallDepth(GenericDatumWriter.class);
+        if (callDepth > 0) {
+          return;
+        }
       AgentSpan span = activeSpan();
       if (span == null) {
         return;