Consult the environment variable when setting the max users frames in code origin probes (#7881)

evanchooly · web-flow · commit 768ce77dbf79 · 2024-11-05T10:04:25.000-05:00
consult the environment variable when setting the max users frames in code origin probes
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java
@@ -100,7 +100,7 @@ public static synchronized void run(
       DebuggerContext.initExceptionDebugger(defaultExceptionDebugger);
     }
     if (config.isDebuggerCodeOriginEnabled()) {
-      DebuggerContext.initCodeOrigin(new DefaultCodeOriginRecorder(configurationUpdater));
+      DebuggerContext.initCodeOrigin(new DefaultCodeOriginRecorder(config, configurationUpdater));
     }
     if (config.isDebuggerInstrumentTheWorld()) {
       setupInstrumentTheWorldTransformer(
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java
@@ -6,6 +6,7 @@
 import com.datadog.debugger.exception.Fingerprinter;
 import com.datadog.debugger.probe.CodeOriginProbe;
 import com.datadog.debugger.probe.Where;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext.CodeOriginRecorder;
@@ -26,6 +27,8 @@
 public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
   private static final Logger LOG = LoggerFactory.getLogger(DefaultCodeOriginRecorder.class);
 
+  private final Config config;
+
   private final ConfigurationUpdater configurationUpdater;
 
   private final Map<String, CodeOriginProbe> fingerprints = new HashMap<>();
@@ -34,7 +37,8 @@ public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
 
   private final AgentTaskScheduler taskScheduler = AgentTaskScheduler.INSTANCE;
 
-  public DefaultCodeOriginRecorder(ConfigurationUpdater configurationUpdater) {
+  public DefaultCodeOriginRecorder(Config config, ConfigurationUpdater configurationUpdater) {
+    this.config = config;
     this.configurationUpdater = configurationUpdater;
   }
 
@@ -59,7 +63,10 @@ public String captureCodeOrigin(String signature) {
 
       probe =
           new CodeOriginProbe(
-              new ProbeId(UUID.randomUUID().toString(), 0), where.getSignature(), where);
+              new ProbeId(UUID.randomUUID().toString(), 0),
+              where.getSignature(),
+              where,
+              config.getDebuggerCodeOriginMaxUserFrames());
       addFingerprint(fingerprint, probe);
 
       installProbe(probe);
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java
@@ -26,14 +26,17 @@
 public class CodeOriginProbe extends LogProbe implements ForceMethodInstrumentation {
   private static final Logger LOGGER = LoggerFactory.getLogger(CodeOriginProbe.class);
 
+  public final int maxFrames;
+
   private final String signature;
 
   private final boolean entrySpanProbe;
 
-  public CodeOriginProbe(ProbeId probeId, String signature, Where where) {
+  public CodeOriginProbe(ProbeId probeId, String signature, Where where, int maxFrames) {
     super(LANGUAGE, probeId, null, where, MethodLocation.EXIT, null, null, true, null, null, null);
     this.signature = signature;
     this.entrySpanProbe = signature != null;
+    this.maxFrames = maxFrames;
   }
 
   @Override
@@ -133,6 +136,7 @@ private List<StackTraceElement> getUserStackFrames() {
         stream ->
             stream
                 .filter(element -> !DebuggerContext.isClassNameExcluded(element.getClassName()))
+                .limit(maxFrames)
                 .collect(Collectors.toList()));
   }
 }
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturingTestBase.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturingTestBase.java
@@ -358,6 +358,7 @@ protected TestSnapshotListener installProbes(
     when(config.getFinalDebuggerSnapshotUrl())
         .thenReturn("http://localhost:8126/debugger/v1/input");
     when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
+    when(config.getDebuggerCodeOriginMaxUserFrames()).thenReturn(20);
     instrumentationListener = new MockInstrumentationListener();
     probeStatusSink = mock(ProbeStatusSink.class);
 
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java
@@ -49,6 +49,8 @@ public class CodeOriginTest extends CapturingTestBase {
   private static final ProbeId CODE_ORIGIN_ID1 = new ProbeId("code origin 1", 0);
   private static final ProbeId CODE_ORIGIN_ID2 = new ProbeId("code origin 2", 0);
 
+  private static final int MAX_FRAMES = 20;
+
   private DefaultCodeOriginRecorder codeOriginRecorder;
 
   private TestTraceInterceptor traceInterceptor;
@@ -112,6 +114,30 @@ public void withLogProbe() throws IOException, URISyntaxException {
     checkResults(testClass, "debug_1", true);
   }
 
+  @Test
+  public void stackDepth() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
+    installProbes(
+        new CodeOriginProbe(
+            CODE_ORIGIN_ID1, null, Where.of(CLASS_NAME, "exit", "()", "39"), MAX_FRAMES));
+
+    Class<?> testClass = compileAndLoadClass("com.datadog.debugger.CodeOrigin04");
+    countFrames(testClass, 10);
+    countFrames(testClass, 100);
+  }
+
+  private void countFrames(Class<?> testClass, int loops) {
+    int result = Reflect.onClass(testClass).call("main", loops).get();
+    assertEquals(loops, result);
+    long count =
+        traceInterceptor.getTrace().stream()
+            .filter(s -> s.getOperationName().equals("exit"))
+            .flatMap(s -> s.getTags().keySet().stream())
+            .filter(key -> key.contains("frames") && key.endsWith("method"))
+            .count();
+    assertTrue(count <= MAX_FRAMES);
+  }
+
   @Test
   public void testCaptureCodeOriginWithSignature() {
     installProbes();
@@ -131,9 +157,9 @@ public void testCaptureCodeOriginWithNullSignature() {
   @NotNull
   private List<LogProbe> codeOriginProbes(String type) {
     CodeOriginProbe entry =
-        new CodeOriginProbe(CODE_ORIGIN_ID1, "()", Where.of(type, "entry", "()", "53"));
+        new CodeOriginProbe(CODE_ORIGIN_ID1, "()", Where.of(type, "entry", "()", "53"), MAX_FRAMES);
     CodeOriginProbe exit =
-        new CodeOriginProbe(CODE_ORIGIN_ID2, null, Where.of(type, "exit", "()", "60"));
+        new CodeOriginProbe(CODE_ORIGIN_ID2, null, Where.of(type, "exit", "()", "60"), MAX_FRAMES);
     return new ArrayList<>(asList(entry, exit));
   }
 
@@ -164,7 +190,7 @@ protected TestSnapshotListener installProbes(LogProbe... probes) {
     TestSnapshotListener listener = super.installProbes(probes);
 
     DebuggerContext.initClassNameFilter(classNameFilter);
-    codeOriginRecorder = new DefaultCodeOriginRecorder(configurationUpdater);
+    codeOriginRecorder = new DefaultCodeOriginRecorder(config, configurationUpdater);
     DebuggerContext.initCodeOrigin(codeOriginRecorder);
 
     return listener;
diff --git a/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CodeOrigin04.java b/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CodeOrigin04.java
@@ -0,0 +1,42 @@
+package com.datadog.debugger;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
+import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
+
+public class CodeOrigin04 {
+  private int intField = 42;
+
+  private static TracerAPI tracerAPI = AgentTracer.get();
+
+  public static int main(int level) throws ReflectiveOperationException {
+    doExit(level);
+
+    return level;
+  }
+
+  private static void doExit(int level) {
+    if (level > 0) {
+      doExit(level - 1);
+    } else {
+      AgentSpan span;
+      AgentScope scope;
+      span = newSpan("exit");
+      scope = tracerAPI.activateSpan(span, ScopeSource.MANUAL);
+      exit();
+      span.finish();
+      scope.close();
+    }
+  }
+
+  private static AgentSpan newSpan(String name) {
+    return tracerAPI.buildSpan("code origin tests", name).start();
+  }
+
+  private static void exit() {
+    int x = 47 / 3;
+  }
+
+}
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -49,6 +49,7 @@ public final class ConfigDefaults {
   static final String DEFAULT_SITE = "datadoghq.com";
 
   static final boolean DEFAULT_CODE_ORIGIN_FOR_SPANS_ENABLED = false;
+  static final int DEFAULT_CODE_ORIGIN_MAX_USER_FRAMES = 8;
   static final boolean DEFAULT_TRACE_SPAN_ORIGIN_ENRICHED = false;
   static final boolean DEFAULT_TRACE_ENABLED = true;
   public static final boolean DEFAULT_TRACE_OTEL_ENABLED = false;
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java
@@ -10,6 +10,7 @@
  */
 public final class TraceInstrumentationConfig {
   public static final String CODE_ORIGIN_FOR_SPANS_ENABLED = "code.origin.for.spans.enabled";
+  public static final String CODE_ORIGIN_MAX_USER_FRAMES = "code.origin.max.user.frames";
   public static final String TRACE_ENABLED = "trace.enabled";
   public static final String TRACE_OTEL_ENABLED = "trace.otel.enabled";
   public static final String INTEGRATIONS_ENABLED = "integrations.enabled";
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -403,6 +403,7 @@ public static String getHostName() {
   private final int debuggerExceptionMaxCapturedFrames;
   private final int debuggerExceptionCaptureInterval;
   private final boolean debuggerCodeOriginEnabled;
+  private final int debuggerCodeOriginMaxUserFrames;
 
   private final Set<String> debuggerThirdPartyIncludes;
   private final Set<String> debuggerThirdPartyExcludes;
@@ -1533,6 +1534,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     debuggerCodeOriginEnabled =
         configProvider.getBoolean(
             CODE_ORIGIN_FOR_SPANS_ENABLED, DEFAULT_CODE_ORIGIN_FOR_SPANS_ENABLED);
+    debuggerCodeOriginMaxUserFrames =
+        configProvider.getInteger(CODE_ORIGIN_MAX_USER_FRAMES, DEFAULT_CODE_ORIGIN_MAX_USER_FRAMES);
     debuggerMaxExceptionPerSecond =
         configProvider.getInteger(
             DEBUGGER_MAX_EXCEPTION_PER_SECOND, DEFAULT_DEBUGGER_MAX_EXCEPTION_PER_SECOND);
@@ -2970,6 +2973,10 @@ public boolean isDebuggerCodeOriginEnabled() {
     return debuggerCodeOriginEnabled;
   }
 
+  public int getDebuggerCodeOriginMaxUserFrames() {
+    return debuggerCodeOriginMaxUserFrames;
+  }
+
   public Set<String> getThirdPartyIncludes() {
     return debuggerThirdPartyIncludes;
   }

Original file line number	Diff line number	Diff line change
`@@ -100,7 +100,7 @@ public static synchronized void run(`
`100`	`100`	`DebuggerContext.initExceptionDebugger(defaultExceptionDebugger);`
`101`	`101`	`}`
`102`	`102`	`if (config.isDebuggerCodeOriginEnabled()) {`
`103`		`- DebuggerContext.initCodeOrigin(new DefaultCodeOriginRecorder(configurationUpdater));`
	`103`	`+ DebuggerContext.initCodeOrigin(new DefaultCodeOriginRecorder(config, configurationUpdater));`
`104`	`104`	`}`
`105`	`105`	`if (config.isDebuggerInstrumentTheWorld()) {`
`106`	`106`	`setupInstrumentTheWorldTransformer(`