Merge pull request #5692 from DataDog/jpbempel/fix-sampling

Fix sampling when log probe is evaluation at Exit
DataDog · Aug 8, 2023 · 5bad76d · 5bad76d
2 parents e195b14 + b760261
commit 5bad76d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 29 deletions.
diff --git a/...ugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java b/...ugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java
@@ -25,6 +25,7 @@
 import static org.objectweb.asm.Type.getType;
 
 import com.datadog.debugger.probe.ProbeDefinition;
+import com.datadog.debugger.probe.SpanDecorationProbe;
 import com.datadog.debugger.probe.Where;
 import com.datadog.debugger.sink.Snapshot;
 import datadog.trace.api.Config;
@@ -302,8 +303,9 @@ private void instrumentMethodEnter() {
       throwableListVar = declareThrowableList(insnList);
     }
     insnList.add(contextInitLabel);
-    if (definition.getEvaluateAt() == MethodLocation.EXIT) {
-      // if evaluation is at exit, skip collecting data at enter
+    if (definition instanceof SpanDecorationProbe
+        && definition.getEvaluateAt() == MethodLocation.EXIT) {
+      // if evaluation is at exit for a span decoration probe, skip collecting data at enter
       methodNode.instructions.insert(methodEnterLabel, insnList);
       return;
     }
@@ -315,32 +317,38 @@ private void instrumentMethodEnter() {
     LabelNode targetNode = new LabelNode();
     LabelNode gotoNode = new LabelNode();
     insnList.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
-    LabelNode inProbeStartLabel = new LabelNode();
-    insnList.add(inProbeStartLabel);
-    // stack []
-    insnList.add(collectCapturedContext(Snapshot.Kind.ENTER, null));
-    // stack [capturedcontext]
-    ldc(insnList, Type.getObjectType(classNode.name));
-    // stack [capturedcontext, class]
-    ldc(insnList, -1L);
-    // stack [capturedcontext, class, long]
-    getStatic(insnList, METHOD_LOCATION_TYPE, "ENTRY");
-    // stack [capturedcontext, class, long, methodlocation]
-    pushProbesIds(insnList);
-    // stack [capturedcontext, class, long, methodlocation, array]
-    invokeStatic(
-        insnList,
-        DEBUGGER_CONTEXT_TYPE,
-        "evalContext",
-        VOID_TYPE,
-        CAPTURED_CONTEXT_TYPE,
-        CLASS_TYPE,
-        LONG_TYPE,
-        METHOD_LOCATION_TYPE,
-        STRING_ARRAY_TYPE);
-    invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "disableInProbe", VOID_TYPE);
-    LabelNode inProbeEndLabel = new LabelNode();
-    insnList.add(inProbeEndLabel);
+    // if evaluation is at exit, skip collecting data at enter
+    if (definition.getEvaluateAt() != MethodLocation.EXIT) {
+      LabelNode inProbeStartLabel = new LabelNode();
+      insnList.add(inProbeStartLabel);
+      // stack []
+      insnList.add(collectCapturedContext(Snapshot.Kind.ENTER, null));
+      // stack [capturedcontext]
+      ldc(insnList, Type.getObjectType(classNode.name));
+      // stack [capturedcontext, class]
+      ldc(insnList, -1L);
+      // stack [capturedcontext, class, long]
+      getStatic(insnList, METHOD_LOCATION_TYPE, "ENTRY");
+      // stack [capturedcontext, class, long, methodlocation]
+      pushProbesIds(insnList);
+      // stack [capturedcontext, class, long, methodlocation, array]
+      invokeStatic(
+          insnList,
+          DEBUGGER_CONTEXT_TYPE,
+          "evalContext",
+          VOID_TYPE,
+          CAPTURED_CONTEXT_TYPE,
+          CLASS_TYPE,
+          LONG_TYPE,
+          METHOD_LOCATION_TYPE,
+          STRING_ARRAY_TYPE);
+      invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "disableInProbe", VOID_TYPE);
+      LabelNode inProbeEndLabel = new LabelNode();
+      insnList.add(inProbeEndLabel);
+      createInProbeFinallyHandler(inProbeStartLabel, inProbeEndLabel);
+    } else {
+      invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "disableInProbe", VOID_TYPE);
+    }
     // stack []
     insnList.add(new JumpInsnNode(Opcodes.GOTO, gotoNode));
     insnList.add(targetNode);
@@ -350,7 +358,6 @@ private void instrumentMethodEnter() {
     // stack []
     insnList.add(gotoNode);
     methodNode.instructions.insert(methodEnterLabel, insnList);
-    createInProbeFinallyHandler(inProbeStartLabel, inProbeEndLabel);
   }
 
   private void createInProbeFinallyHandler(LabelNode inProbeStartLabel, LabelNode inProbeEndLabel) {

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
@@ -1828,6 +1828,7 @@ private static LogProbe.Builder createProbeBuilder(
         .probeId(id)
         .captureSnapshot(true)
         .where(typeName, methodName, signature, lines)
+        // Increase sampling limit to avoid being sampled during tests
         .sampling(new LogProbe.Sampling(100));
   }