Add evaluateAt support for probe instrumentations

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/Snapshot.java

@@ -92,7 +92,9 @@ private void addCapturingProbeId(ProbeDetails probe) {
   public void setEntry(CapturedContext context) {
     summaryBuilder.addEntry(context);
     context.setThisClassName(thisClassName);
-    if (checkCapture(context)) {
+    if ((probe.getEvaluateAt() == MethodLocation.DEFAULT
+            || probe.getEvaluateAt() == MethodLocation.ENTRY)
+        && checkCapture(context)) {
       captures.setEntry(context);
     }
   }
@@ -102,7 +104,9 @@ public void setExit(CapturedContext context) {
     context.addExtension(ValueReferences.DURATION_EXTENSION_NAME, duration);
     summaryBuilder.addExit(context);
     context.setThisClassName(thisClassName);
-    if (checkCapture(context)) {
+    if ((probe.getEvaluateAt() == MethodLocation.DEFAULT
+            || probe.getEvaluateAt() == MethodLocation.EXIT)
+        && checkCapture(context)) {
       captures.setReturn(context);
     }
   }
@@ -220,7 +224,8 @@ private Snapshot copy(String probeId, String newSnapshotId) {
         duration,
         stack,
         captures,
-        new ProbeDetails(probeId, probe.location, probe.script, probe.tags, summaryBuilder),
+        new ProbeDetails(
+            probeId, probe.location, probe.evaluateAt, probe.script, probe.tags, summaryBuilder),
         language,
         thread,
         thisClassName,
@@ -290,6 +295,12 @@ public enum Kind {
     AFTER;
   }
 
+  public enum MethodLocation {
+    DEFAULT,
+    ENTRY,
+    EXIT
+  }
+
   /** Probe information associated with a snapshot */
   public static class ProbeDetails {
     public static final String ITW_PROBE_ID = "instrument-the-world-probe";
@@ -299,33 +310,44 @@ public static class ProbeDetails {
 
     private final String id;
     private final ProbeLocation location;
+    private final MethodLocation evaluateAt;
     private final DebuggerScript script;
     private final transient List<ProbeDetails> additionalProbes;
     private final String tags;
     private final transient SummaryBuilder summaryBuilder;
 
     public ProbeDetails(String id, ProbeLocation location) {
-      this(id, location, null, null, new SnapshotSummaryBuilder(location), Collections.emptyList());
+      this(
+          id,
+          location,
+          MethodLocation.DEFAULT,
+          null,
+          null,
+          new SnapshotSummaryBuilder(location),
+          Collections.emptyList());
     }
 
     public ProbeDetails(
         String id,
         ProbeLocation location,
+        MethodLocation evaluateAt,
         DebuggerScript script,
         String tags,
         SummaryBuilder summaryBuilder) {
-      this(id, location, script, tags, summaryBuilder, Collections.emptyList());
+      this(id, location, evaluateAt, script, tags, summaryBuilder, Collections.emptyList());
     }
 
     public ProbeDetails(
         String id,
         ProbeLocation location,
+        MethodLocation evaluateAt,
         DebuggerScript script,
         String tags,
         SummaryBuilder summaryBuilder,
         List<ProbeDetails> additionalProbes) {
       this.id = id;
       this.location = location;
+      this.evaluateAt = evaluateAt;
       this.script = script;
       this.additionalProbes = additionalProbes;
       this.tags = tags;
@@ -340,6 +362,10 @@ public ProbeLocation getLocation() {
       return location;
     }
 
+    public MethodLocation getEvaluateAt() {
+      return evaluateAt;
+    }
+
     public DebuggerScript getScript() {
       return script;
     }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/ConfigurationUpdater.java

@@ -294,6 +294,7 @@ private Snapshot.ProbeDetails convertToProbeDetails(
     return new Snapshot.ProbeDetails(
         probe.getId(),
         location,
+        convertMethodLocation(probe.getEvaluateAt()),
         probeCondition,
         probe.concatTags(),
         summaryBuilder,
@@ -302,6 +303,19 @@ private Snapshot.ProbeDetails convertToProbeDetails(
             .collect(Collectors.toList()));
   }
 
+  private Snapshot.MethodLocation convertMethodLocation(
+      ProbeDefinition.MethodLocation methodLocation) {
+    switch (methodLocation) {
+      case DEFAULT:
+        return Snapshot.MethodLocation.DEFAULT;
+      case ENTRY:
+        return Snapshot.MethodLocation.ENTRY;
+      case EXIT:
+        return Snapshot.MethodLocation.EXIT;
+    }
+    return null;
+  }
+
   private void applyRateLimiter(ConfigurationComparer changes) {
     Collection<SnapshotProbe> probes = currentConfiguration.getSnapshotProbes();
     if (probes == null) {

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

@@ -36,6 +36,7 @@ public class Instrumentor {
   protected int localVarBaseOffset;
   protected int argOffset;
   protected final String[] argumentNames;
+  protected LabelNode returnHandlerLabel;
 
   public Instrumentor(
       ProbeDefinition definition,
@@ -134,6 +135,70 @@ protected void fillLineMap() {
     }
   }
 
+  protected void processInstructions() {
+    AbstractInsnNode node = methodNode.instructions.getFirst();
+    while (node != null && !node.equals(returnHandlerLabel)) {
+      if (node.getType() == AbstractInsnNode.LINE) {
+        lineMap.addLine((LineNumberNode) node);
+      } else {
+        node = processInstruction(node);
+      }
+      node = node.getNext();
+    }
+    if (returnHandlerLabel == null) {
+      // if no return found, fallback to use the last instruction as last resort
+      returnHandlerLabel = new LabelNode();
+      methodNode.instructions.insert(methodNode.instructions.getLast(), returnHandlerLabel);
+    }
+  }
+
+  protected AbstractInsnNode processInstruction(AbstractInsnNode node) {
+    switch (node.getOpcode()) {
+      case Opcodes.RET:
+      case Opcodes.RETURN:
+      case Opcodes.IRETURN:
+      case Opcodes.FRETURN:
+      case Opcodes.LRETURN:
+      case Opcodes.DRETURN:
+      case Opcodes.ARETURN:
+        {
+          InsnList beforeReturnInsnList = getBeforeReturnInsnList(node);
+          methodNode.instructions.insertBefore(node, beforeReturnInsnList);
+          AbstractInsnNode prev = node.getPrevious();
+          methodNode.instructions.remove(node);
+          methodNode.instructions.insert(
+              prev, new JumpInsnNode(Opcodes.GOTO, getReturnHandler(node)));
+          return prev;
+        }
+    }
+    return node;
+  }
+
+  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+    return null;
+  }
+
+  protected LabelNode getReturnHandler(AbstractInsnNode exitNode) {
+    // exit node must have been removed from the original instruction list
+    if (exitNode.getNext() != null || exitNode.getPrevious() != null) {
+      throw new IllegalArgumentException("exitNode is not removed from original instruction list");
+    }
+    if (returnHandlerLabel != null) {
+      return returnHandlerLabel;
+    }
+    returnHandlerLabel = new LabelNode();
+    methodNode.instructions.add(returnHandlerLabel);
+    // stack top is return value (if any)
+    InsnList handler = getReturnHandlerInsnList();
+    handler.add(exitNode); // stack: []
+    methodNode.instructions.add(handler);
+    return returnHandlerLabel;
+  }
+
+  protected InsnList getReturnHandlerInsnList() {
+    return new InsnList();
+  }
+
   protected static void invokeStatic(
       InsnList insnList, Type owner, String name, Type returnType, Type... argTypes) {
     // expected stack: [arg_type_1 ... arg_type_N]

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

@@ -17,6 +17,7 @@
 import java.util.regex.Pattern;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.ClassNode;
 import org.objectweb.asm.tree.FieldInsnNode;
 import org.objectweb.asm.tree.FieldNode;
@@ -55,11 +56,31 @@ public void instrument() {
       fillLineMap();
       addLineMetric(lineMap);
     } else {
-      InsnList insnList = callMetric(metricProbe);
-      methodNode.instructions.insert(methodEnterLabel, insnList);
+      switch (definition.getEvaluateAt()) {
+        case ENTRY:
+        case DEFAULT:
+          {
+            InsnList insnList = callMetric(metricProbe);
+            methodNode.instructions.insert(methodEnterLabel, insnList);
+            break;
+          }
+        case EXIT:
+          {
+            processInstructions();
+            break;
+          }
+        default:
+          throw new IllegalArgumentException(
+              "Invalid evaluateAt attribute: " + definition.getEvaluateAt());
+      }
     }
   }
 
+  @Override
+  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+    return callMetric(metricProbe);
+  }
+
   private InsnList callCount(MetricProbe metricProbe) {
     InsnList insnList = new InsnList();
     if (metricProbe.getValue() == null) {

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

@@ -12,6 +12,7 @@
 import static org.objectweb.asm.Type.INT_TYPE;
 
 import com.datadog.debugger.probe.LogProbe;
+import com.datadog.debugger.probe.ProbeDefinition;
 import com.datadog.debugger.probe.SnapshotProbe;
 import com.datadog.debugger.probe.Where;
 import datadog.trace.api.Config;
@@ -31,7 +32,6 @@
 import org.objectweb.asm.tree.InsnNode;
 import org.objectweb.asm.tree.JumpInsnNode;
 import org.objectweb.asm.tree.LabelNode;
-import org.objectweb.asm.tree.LineNumberNode;
 import org.objectweb.asm.tree.LocalVariableNode;
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
@@ -45,7 +45,6 @@ public final class SnapshotInstrumentor extends Instrumentor {
   private final boolean captureFullState;
   private final LabelNode snapshotInitLabel = new LabelNode();
   private int snapshotVar = -1;
-  private LabelNode returnHandlerLabel = null;
 
   public SnapshotInstrumentor(
       SnapshotProbe snapshotProbe,
@@ -81,45 +80,6 @@ public void instrument() {
     }
   }
 
-  private void processInstructions() {
-    AbstractInsnNode node = methodNode.instructions.getFirst();
-    while (node != null && !node.equals(returnHandlerLabel)) {
-      if (node.getType() == AbstractInsnNode.LINE) {
-        lineMap.addLine((LineNumberNode) node);
-      } else {
-        node = processInstruction(node);
-      }
-      node = node.getNext();
-    }
-    if (returnHandlerLabel == null) {
-      // if no return found, fallback to use the last instruction as last resort
-      returnHandlerLabel = new LabelNode();
-      methodNode.instructions.insert(methodNode.instructions.getLast(), returnHandlerLabel);
-    }
-  }
-
-  private AbstractInsnNode processInstruction(AbstractInsnNode node) {
-    switch (node.getOpcode()) {
-      case Opcodes.RET:
-      case Opcodes.RETURN:
-      case Opcodes.IRETURN:
-      case Opcodes.FRETURN:
-      case Opcodes.LRETURN:
-      case Opcodes.DRETURN:
-      case Opcodes.ARETURN:
-        {
-          methodNode.instructions.insertBefore(
-              node, collectSnapshotCapture(-1, Snapshot.Kind.RETURN, node));
-          AbstractInsnNode prev = node.getPrevious();
-          methodNode.instructions.remove(node);
-          methodNode.instructions.insert(
-              prev, new JumpInsnNode(Opcodes.GOTO, getReturnHandler(node)));
-          return prev;
-        }
-    }
-    return node;
-  }
-
   private void addLineCaptures(LineMap lineMap) {
     Where.SourceLine[] targetLines = definition.getWhere().getSourceLines();
     if (targetLines == null) {
@@ -156,19 +116,14 @@ private void addLineCaptures(LineMap lineMap) {
     }
   }
 
-  private LabelNode getReturnHandler(AbstractInsnNode exitNode) {
-    // exit node must have been removed from the original instruction list
-    assert exitNode.getNext() == null && exitNode.getPrevious() == null;
-    if (returnHandlerLabel != null) {
-      return returnHandlerLabel;
-    }
-    returnHandlerLabel = new LabelNode();
-    methodNode.instructions.add(returnHandlerLabel);
-    // stack top is return value (if any)
-    InsnList handler = commitSnapshot();
-    handler.add(exitNode); // stack: []
-    methodNode.instructions.add(handler);
-    return returnHandlerLabel;
+  @Override
+  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+    return collectSnapshotCapture(-1, Snapshot.Kind.RETURN, node);
+  }
+
+  @Override
+  protected InsnList getReturnHandlerInsnList() {
+    return commitSnapshot();
   }
 
   private InsnList commitSnapshot() {
@@ -203,8 +158,16 @@ private void addFinallyHandler(LabelNode endLabel) {
   }
 
   private void instrumentMethodEnter() {
-    methodNode.instructions.insert(
-        snapshotInitLabel, collectSnapshotCapture(-1, Snapshot.Kind.ENTER, null));
+    InsnList insnList;
+    if (definition.getEvaluateAt() == ProbeDefinition.MethodLocation.EXIT) {
+      // if evaluation is at exit, skip collecting data at enter
+      // but create the snapshot at entry anyway
+      insnList = new InsnList();
+      getSnapshot(insnList);
+    } else {
+      insnList = collectSnapshotCapture(-1, Snapshot.Kind.ENTER, null);
+    }
+    methodNode.instructions.insert(snapshotInitLabel, insnList);
   }
 
   private void instrumentTryCatchHandlers() {

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

@@ -85,7 +85,7 @@ public String toString() {
   // no-arg constructor is required by Moshi to avoid creating instance with unsafe and by-passing
   // constructors, including field initializers.
   public LogProbe() {
-    this(LANGUAGE, null, true, null, null, MethodLocation.NONE, null, new ArrayList<>());
+    this(LANGUAGE, null, true, null, null, MethodLocation.DEFAULT, null, new ArrayList<>());
   }
 
   public LogProbe(
@@ -180,7 +180,7 @@ public static class Builder extends ProbeDefinition.Builder<Builder> {
     private String template;
     private List<Segment> segments;
 
-    public LogProbe.Builder template(String template) {
+    public Builder template(String template) {
       this.template = template;
       this.segments = parseTemplate(template);
       return this;

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

@@ -26,7 +26,7 @@ public enum MetricKind {
   // no-arg constructor is required by Moshi to avoid creating instance with unsafe and by-passing
   // constructors, including field initializers.
   public MetricProbe() {
-    this(LANGUAGE, null, true, null, null, MethodLocation.NONE, MetricKind.COUNT, null, null);
+    this(LANGUAGE, null, true, null, null, MethodLocation.DEFAULT, MetricKind.COUNT, null, null);
   }
 
   public MetricProbe(