Capture values at entry for method probe

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)

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) {

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 []

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;

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();
   }
 

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();
   }
 

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();
   }
 

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