Capture values at entry for method probe

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.

Author: jpbempel

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;
@@ -691,6 +692,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) {
@@ -702,6 +704,10 @@ private ProbeDefinition selectReferenceDefinition(
         LogProbe logProbe = (LogProbe) definition;
         captureSnapshot = captureSnapshot | logProbe.isCaptureSnapshot();
         capture = mergeCapture(capture, logProbe.getCapture());
+        probeCondition =
+            probeCondition == null && logProbe.getProbeCondition() != null
+                ? logProbe.getProbeCondition()
+                : probeCondition;
       }
       if (definition.getEvaluateAt() == MethodLocation.ENTRY
           || definition.getEvaluateAt() == MethodLocation.DEFAULT) {
@@ -711,6 +717,7 @@ private ProbeDefinition selectReferenceDefinition(
     if (hasLogProbe) {
       return LogProbe.builder()
           .probeId(probeId)
+          .when(probeCondition)
           .where(where)
           .evaluateAt(evaluateAt)
           .capture(capture)

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java

@@ -424,8 +424,8 @@ 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 evaluation is at exit and with condition, skip collecting data at enter
+    if (definition.getEvaluateAt() != MethodLocation.EXIT || !definition.hasCondition()) {
       LabelNode inProbeStartLabel = new LabelNode();
       insnList.add(inProbeStartLabel);
       // stack []

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ExceptionProbe.java

@@ -61,12 +61,18 @@ public CapturedContext.Status createStatus() {
   @Override
   public void evaluate(
       CapturedContext context, CapturedContext.Status status, MethodLocation methodLocation) {
-    if (!(status instanceof ExceptionProbeStatus)) {
+    ExceptionProbeStatus exceptionStatus;
+    if (status instanceof ExceptionProbeStatus) {
+      exceptionStatus = (ExceptionProbeStatus) status;
+    } else {
       throw new IllegalStateException("Invalid status: " + status.getClass());
     }
     if (methodLocation != MethodLocation.EXIT) {
       return;
     }
+    // reset capture flag as we are at exit, consider at entry always capture
+    // only the evaluation at exit decide to send the snapshot or not based on exception evaluation
+    exceptionStatus.setCapture(false);
     if (context.getCapturedThrowable() == null) {
       return;
     }
@@ -81,7 +87,6 @@ public void evaluate(
     if (exceptionProbeManager.shouldCaptureException(fingerprint)) {
       LOGGER.debug("Capturing exception matching fingerprint: {}", fingerprint);
       // capture only on uncaught exception matching the fingerprint
-      ExceptionProbeStatus exceptionStatus = (ExceptionProbeStatus) status;
       ExceptionProbeManager.ThrowableState state =
           exceptionProbeManager.getStateByThrowable(innerMostThrowable);
       if (state != null) {
@@ -147,7 +152,7 @@ public void buildLocation(InstrumentationResult result) {
   }
 
   public static class ExceptionProbeStatus extends LogStatus {
-    private boolean capture;
+    private boolean capture = true;
 
     public ExceptionProbeStatus(ProbeImplementation probeImplementation) {
       super(probeImplementation);

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";