Capture values at entry for method probe (#8169)

jpbempel · web-flow · commit 8b63e5a85b23 · 2025-01-14T22:34:52.000+01:00
method probe when evaluateAt is Exit and no condition should capture
entry values. When there is a condition, entry values are not captured
instrumentation check the presence of the condition to generate the
entry instrumentation or not. therefore, the dummy definition used
when duplicate probes need to merge the probe conditions if at least
one probe has one condition.
We are delegating the decision to capture through a new field in
CapturedContextInstrumentor
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
@@ -3,6 +3,7 @@
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 
+import com.datadog.debugger.el.ProbeCondition;
 import com.datadog.debugger.instrumentation.DiagnosticMessage;
 import com.datadog.debugger.instrumentation.InstrumentationResult;
 import com.datadog.debugger.instrumentation.MethodInfo;
@@ -693,6 +694,7 @@ private ProbeDefinition selectReferenceDefinition(
     MethodLocation evaluateAt = MethodLocation.EXIT;
     LogProbe.Capture capture = null;
     boolean captureSnapshot = false;
+    ProbeCondition probeCondition = null;
     Where where = capturedContextProbes.get(0).getWhere();
     ProbeId probeId = capturedContextProbes.get(0).getProbeId();
     for (ProbeDefinition definition : capturedContextProbes) {
@@ -704,6 +706,9 @@ private ProbeDefinition selectReferenceDefinition(
         LogProbe logProbe = (LogProbe) definition;
         captureSnapshot = captureSnapshot | logProbe.isCaptureSnapshot();
         capture = mergeCapture(capture, logProbe.getCapture());
+        if (probeCondition == null) {
+          probeCondition = logProbe.getProbeCondition();
+        }
       }
       if (definition.getEvaluateAt() == MethodLocation.ENTRY
           || definition.getEvaluateAt() == MethodLocation.DEFAULT) {
@@ -713,6 +718,7 @@ private ProbeDefinition selectReferenceDefinition(
     if (hasLogProbe) {
       return LogProbe.builder()
           .probeId(probeId)
+          .when(probeCondition)
           .where(where)
           .evaluateAt(evaluateAt)
           .capture(capture)
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/ExceptionProbeManager.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/ExceptionProbeManager.java
@@ -118,13 +118,7 @@ private static ExceptionProbe createMethodProbe(
       ExceptionProbeManager exceptionProbeManager, Where where, int chainedExceptionIdx) {
     String probeId = UUID.randomUUID().toString();
     return new ExceptionProbe(
-        new ProbeId(probeId, 0),
-        where,
-        null,
-        null,
-        null,
-        exceptionProbeManager,
-        chainedExceptionIdx);
+        new ProbeId(probeId, 0), where, null, null, exceptionProbeManager, chainedExceptionIdx);
   }
 
   public boolean isAlreadyInstrumented(String fingerprint) {
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java
@@ -69,6 +69,7 @@
 public class CapturedContextInstrumentor extends Instrumentor {
   private static final Logger LOGGER = LoggerFactory.getLogger(CapturedContextInstrumentor.class);
   private final boolean captureSnapshot;
+  private final boolean captureEntry;
   private final Limits limits;
   private final LabelNode contextInitLabel = new LabelNode();
   private int entryContextVar = -1;
@@ -83,9 +84,11 @@ public CapturedContextInstrumentor(
       List<DiagnosticMessage> diagnostics,
       List<ProbeId> probeIds,
       boolean captureSnapshot,
+      boolean captureEntry,
       Limits limits) {
     super(definition, methodInfo, diagnostics, probeIds);
     this.captureSnapshot = captureSnapshot;
+    this.captureEntry = captureEntry;
     this.limits = limits;
   }
 
@@ -422,8 +425,7 @@ private void instrumentMethodEnter() {
     LabelNode targetNode = new LabelNode();
     LabelNode gotoNode = new LabelNode();
     insnList.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
-    // if evaluation is at exit, skip collecting data at enter
-    if (definition.getEvaluateAt() != MethodLocation.EXIT) {
+    if (captureEntry) {
       LabelNode inProbeStartLabel = new LabelNode();
       insnList.add(inProbeStartLabel);
       // stack []
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ExceptionProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ExceptionProbe.java
@@ -3,6 +3,7 @@
 import static com.datadog.debugger.util.ExceptionHelper.getInnerMostThrowable;
 import static java.util.Collections.emptyList;
 
+import com.datadog.debugger.el.DSL;
 import com.datadog.debugger.el.ProbeCondition;
 import com.datadog.debugger.exception.ExceptionProbeManager;
 import com.datadog.debugger.exception.Fingerprinter;
@@ -26,7 +27,6 @@ public class ExceptionProbe extends LogProbe implements ForceMethodInstrumentati
   public ExceptionProbe(
       ProbeId probeId,
       Where where,
-      ProbeCondition probeCondition,
       Capture capture,
       Sampling sampling,
       ExceptionProbeManager exceptionProbeManager,
@@ -40,7 +40,8 @@ public ExceptionProbe(
         null,
         null,
         true,
-        probeCondition,
+        // forcing a useless condition to be instrumented with captureEntry=false
+        new ProbeCondition(DSL.when(DSL.TRUE), "true"),
         capture,
         sampling);
     this.exceptionProbeManager = exceptionProbeManager;
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java
@@ -409,8 +409,16 @@ public Sampling getSampling() {
   @Override
   public InstrumentationResult.Status instrument(
       MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
+    // if evaluation is at exit and with condition, skip collecting data at entry
+    boolean captureEntry = !(getEvaluateAt() == MethodLocation.EXIT && hasCondition());
     return new CapturedContextInstrumentor(
-            this, methodInfo, diagnostics, probeIds, isCaptureSnapshot(), toLimits(getCapture()))
+            this,
+            methodInfo,
+            diagnostics,
+            probeIds,
+            isCaptureSnapshot(),
+            captureEntry,
+            toLimits(getCapture()))
         .instrument();
   }
 
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java
@@ -135,8 +135,9 @@ public SpanDecorationProbe(
   @Override
   public InstrumentationResult.Status instrument(
       MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
+    boolean captureEntry = evaluateAt != MethodLocation.EXIT;
     return new CapturedContextInstrumentor(
-            this, methodInfo, diagnostics, probeIds, false, Limits.DEFAULT)
+            this, methodInfo, diagnostics, probeIds, false, captureEntry, Limits.DEFAULT)
         .instrument();
   }
 
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/TriggerProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/TriggerProbe.java
@@ -54,7 +54,8 @@ public TriggerProbe(ProbeId probeId, Where where) {
   @Override
   public InstrumentationResult.Status instrument(
       MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
-    return new CapturedContextInstrumentor(this, methodInfo, diagnostics, probeIds, false, null)
+    return new CapturedContextInstrumentor(
+            this, methodInfo, diagnostics, probeIds, false, false, null)
         .instrument();
   }
 
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
@@ -134,6 +134,43 @@ public void methodProbe() throws IOException, URISyntaxException {
     assertEquals("CapturedSnapshot01.main", snapshot.getStack().get(0).getFunction());
   }
 
+  @Test
+  public void methodProbeAtExit() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "CapturedSnapshot01";
+    TestSnapshotListener listener =
+        installSingleProbeAtExit(CLASS_NAME, "main", "int (java.lang.String)");
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "1").get();
+    assertEquals(3, result);
+    Snapshot snapshot = assertOneSnapshot(listener);
+    assertCaptureArgs(snapshot.getCaptures().getEntry(), "arg", "java.lang.String", "1");
+    assertCaptureArgs(snapshot.getCaptures().getReturn(), "arg", "java.lang.String", "1");
+    assertTrue(snapshot.getDuration() > 0);
+    assertTrue(snapshot.getStack().size() > 0);
+    assertEquals("CapturedSnapshot01.main", snapshot.getStack().get(0).getFunction());
+  }
+
+  @Test
+  public void methodProbeAtExitWithCondition() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "CapturedSnapshot01";
+    LogProbe probe =
+        createProbeBuilder(PROBE_ID, CLASS_NAME, "main", "int (java.lang.String)")
+            .when(
+                new ProbeCondition(DSL.when(DSL.eq(DSL.ref("arg"), DSL.value("1"))), "arg == '1'"))
+            .evaluateAt(MethodLocation.EXIT)
+            .build();
+    TestSnapshotListener listener = installProbes(probe);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "1").get();
+    assertEquals(3, result);
+    Snapshot snapshot = assertOneSnapshot(listener);
+    assertEquals(CapturedContext.EMPTY_CAPTURING_CONTEXT, snapshot.getCaptures().getEntry());
+    assertCaptureArgs(snapshot.getCaptures().getReturn(), "arg", "java.lang.String", "1");
+    assertTrue(snapshot.getDuration() > 0);
+    assertTrue(snapshot.getStack().size() > 0);
+    assertEquals("CapturedSnapshot01.main", snapshot.getStack().get(0).getFunction());
+  }
+
   @Test
   public void localVarHoistingNoPreviousStore() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.fasterxml.jackson.core.json.ByteSourceJsonBootstrapper";

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,8 @@ public TriggerProbe(ProbeId probeId, Where where) {`
`54`	`54`	`@Override`
`55`	`55`	`public InstrumentationResult.Status instrument(`
`56`	`56`	`MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {`
`57`		`- return new CapturedContextInstrumentor(this, methodInfo, diagnostics, probeIds, false, null)`
	`57`	`+ return new CapturedContextInstrumentor(`
	`58`	`+ this, methodInfo, diagnostics, probeIds, false, false, null)`
`58`	`59`	`.instrument();`
`59`	`60`	`}`
`60`	`61`