Merge pull request #6135 from DataDog/jpbempel/line-probe-rate-limit

Fix sampling for line probe without condition
DataDog · Nov 2, 2023 · 00c9ebe · 00c9ebe
2 parents 84421f3 + ef53842
commit 00c9ebe
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 12 deletions.
diff --git a/...er/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/CapturedContext.java b/...er/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/CapturedContext.java
@@ -327,22 +327,14 @@ public Status evaluate(
           ValueReferences.DURATION_EXTENSION_NAME, duration / 1_000_000.0); // convert to ms
     }
     this.thisClassName = thisClassName;
-    boolean shouldEvaluate = resolveEvaluateAt(probeImplementation, methodLocation);
+    boolean shouldEvaluate =
+        MethodLocation.isSame(methodLocation, probeImplementation.getEvaluateAt());
     if (shouldEvaluate) {
       probeImplementation.evaluate(this, status, methodLocation);
     }
     return status;
   }
 
-  private static boolean resolveEvaluateAt(
-      ProbeImplementation probeImplementation, MethodLocation methodLocation) {
-    if (methodLocation == MethodLocation.DEFAULT) {
-      // line probe, no evaluation of probe's evaluateAt
-      return true;
-    }
-    return MethodLocation.isSame(methodLocation, probeImplementation.getEvaluateAt());
-  }
-
   public Status getStatus(String probeId) {
     Status result = statusByProbeId.get(probeId);
     if (result == null) {

diff --git a/...ger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/MethodLocation.java b/...ger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/MethodLocation.java
@@ -6,6 +6,11 @@ public enum MethodLocation {
   EXIT;
 
   public static boolean isSame(MethodLocation methodLocation, MethodLocation evaluateAt) {
+    if (methodLocation == MethodLocation.DEFAULT) {
+      // line probe, no evaluation of probe's evaluateAt
+      // MethodLocation.DEFAULT is used for line probe when evaluating the context
+      return true;
+    }
     if (methodLocation == MethodLocation.ENTRY) {
       return evaluateAt == MethodLocation.DEFAULT || evaluateAt == MethodLocation.ENTRY;
     }

diff --git a/...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
@@ -148,7 +148,7 @@ public void methodProbe() throws IOException, URISyntaxException {
   public void singleLineProbe() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot01";
     DebuggerTransformerTest.TestSnapshotListener listener =
-        installSingleProbe(CLASS_NAME, "main", "int (java.lang.String)", "8");
+        installSingleProbeAtExit(CLASS_NAME, "main", "int (java.lang.String)", "8");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "1").get();
     assertEquals(3, result);
@@ -1582,7 +1582,7 @@ private Snapshot doUnknownCount(String CLASS_NAME) throws IOException, URISyntax
   public void beforeForLoopLineProbe() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot02";
     DebuggerTransformerTest.TestSnapshotListener listener =
-        installSingleProbe(CLASS_NAME, null, null, "46");
+        installSingleProbeAtExit(CLASS_NAME, null, null, "46");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "synchronizedBlock").get();
     assertEquals(76, result);
@@ -1597,10 +1597,12 @@ public void dupLineProbeSameTemplate() throws IOException, URISyntaxException {
     LogProbe probe1 =
         createProbeBuilder(PROBE_ID1, CLASS_NAME, null, null, "39")
             .template(LOG_TEMPLATE, parseTemplate(LOG_TEMPLATE))
+            .evaluateAt(MethodLocation.EXIT)
             .build();
     LogProbe probe2 =
         createProbeBuilder(PROBE_ID2, CLASS_NAME, null, null, "39")
             .template(LOG_TEMPLATE, parseTemplate(LOG_TEMPLATE))
+            .evaluateAt(MethodLocation.EXIT)
             .build();
     DebuggerTransformerTest.TestSnapshotListener listener =
         installProbes(CLASS_NAME, probe1, probe2);
@@ -1910,6 +1912,12 @@ private DebuggerTransformerTest.TestSnapshotListener installSingleProbe(
     return installProbes(typeName, logProbes);
   }
 
+  private DebuggerTransformerTest.TestSnapshotListener installSingleProbeAtExit(
+      String typeName, String methodName, String signature, String... lines) {
+    LogProbe logProbes = createProbeAtExit(PROBE_ID, typeName, methodName, signature, lines);
+    return installProbes(typeName, logProbes);
+  }
+
   private DebuggerTransformerTest.TestSnapshotListener installProbes(
       String expectedClassName, Configuration configuration) {
     Config config = mock(Config.class);
@@ -2169,6 +2177,13 @@ private static LogProbe createProbe(
     return createProbeBuilder(id, typeName, methodName, signature, lines).build();
   }
 
+  private static LogProbe createProbeAtExit(
+      ProbeId id, String typeName, String methodName, String signature, String... lines) {
+    return createProbeBuilder(id, typeName, methodName, signature, lines)
+        .evaluateAt(MethodLocation.EXIT)
+        .build();
+  }
+
   private static LogProbe.Builder createProbeBuilder(
       ProbeId id, String typeName, String methodName, String signature, String... lines) {
     return LogProbe.builder()
@@ -2186,6 +2201,7 @@ private static LogProbe createSourceFileProbe(ProbeId id, String sourceFile, int
         .probeId(id)
         .captureSnapshot(true)
         .where(null, null, null, line, sourceFile)
+        .evaluateAt(MethodLocation.EXIT)
         .build();
   }