Fix instrumenting metric probes for Kotlin (#11890)

Fix instrumenting metric probes for Kotlin

track objects on the stack with frame for each instruction using
ASM analyzer. based on this generate POP/POP2 bytecode instruction
to leave the stack clean before inserting metric calls.

Kotlin compiler can, especially for suspend functions, generates
a state machine with duplicate values on the stack while a return
instruction will clear them (dropping the frame).

we are previously relying on the fact that the number of values on the
stack are precise even before a return instruction which leads to
introduce jump to a same location but without the same number of
values on the stack, triggering the verifier.

fix leftover

Co-authored-by: jean-philippe.bempel <jean-philippe.bempel@datadoghq.com>

Author: jpbempel

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

@@ -40,6 +40,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.objectweb.asm.Opcodes;
@@ -58,6 +59,8 @@
 import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.TypeInsnNode;
 import org.objectweb.asm.tree.VarInsnNode;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -265,7 +268,8 @@ private boolean isExceptionLocalDeclared(TryCatchBlockNode catchHandler, MethodN
   }
 
   @Override
-  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+  protected InsnList getBeforeReturnInsnList(
+      AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     InsnList insnList = new InsnList();
     // stack [ret_value]
     insnList.add(new VarInsnNode(Opcodes.ALOAD, entryContextVar));

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

@@ -3,8 +3,11 @@
 import com.datadog.debugger.probe.ProbeDefinition;
 import datadog.trace.bootstrap.debugger.Limits;
 import java.util.List;
+import java.util.Map;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
 
 public class ExceptionInstrumenter extends CapturedContextInstrumenter {
 
@@ -25,7 +28,8 @@ public InstrumentationResult.Status instrument() {
   }
 
   @Override
-  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+  protected InsnList getBeforeReturnInsnList(
+      AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     return null;
   }
 

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

@@ -13,6 +13,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import org.objectweb.asm.Opcodes;
@@ -28,9 +29,17 @@
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.TypeInsnNode;
+import org.objectweb.asm.tree.analysis.Analyzer;
+import org.objectweb.asm.tree.analysis.AnalyzerException;
+import org.objectweb.asm.tree.analysis.BasicInterpreter;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** Common class for generating instrumentation */
 public abstract class Instrumenter {
+  private static final Logger LOGGER = LoggerFactory.getLogger(Instrumenter.class);
   protected static final String CONSTRUCTOR_NAME = "<init>";
   protected static final String PROBEID_TAG_NAME = "debugger.probeid";
 
@@ -151,12 +160,14 @@ private AbstractInsnNode findFirstInsnForConstructor(AbstractInsnNode first) {
   }
 
   protected void processInstructions() {
+    Map<AbstractInsnNode, Frame<BasicValue>> frames = new IdentityHashMap<>();
+    computeFrames(classNode.name, methodNode, frames);
     AbstractInsnNode node = methodNode.instructions.getFirst();
     LabelNode sentinelNode = new LabelNode();
     methodNode.instructions.add(sentinelNode);
     while (node != null && !node.equals(sentinelNode)) {
       if (node.getType() != AbstractInsnNode.LINE) {
-        node = processInstruction(node);
+        node = processInstruction(node, frames);
       }
       node = node.getNext();
     }
@@ -168,7 +179,40 @@ protected void processInstructions() {
     }
   }
 
-  protected AbstractInsnNode processInstruction(AbstractInsnNode node) {
+  // returns the extra values sitting *below* `valuesToKeep` values already
+  // accounted for at the top of the stack (e.g. the return value), as POP/POP2 insns
+  protected InsnList stackCleanupInsnList(
+      AbstractInsnNode node, int valuesToKeep, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
+    InsnList result = new InsnList();
+    Frame<BasicValue> frame = frames.get(node);
+    if (frame == null) {
+      return result;
+    }
+    for (int i = frame.getStackSize() - valuesToKeep - 1; i >= 0; i--) {
+      BasicValue value = frame.getStack(i);
+      result.add(new InsnNode(value.getSize() == 2 ? Opcodes.POP2 : Opcodes.POP));
+    }
+    return result;
+  }
+
+  private static void computeFrames(
+      String owner, MethodNode methodNode, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
+    try {
+      Frame<BasicValue>[] frameArray =
+          new Analyzer<>(new BasicInterpreter()).analyze(owner, methodNode);
+      AbstractInsnNode current = methodNode.instructions.getFirst();
+      int idx = 0;
+      while (current != null) {
+        frames.put(current, frameArray[idx++]);
+        current = current.getNext();
+      }
+    } catch (AnalyzerException ex) {
+      LOGGER.debug("Failed to analyze method[%s::%s] instructions", owner, methodNode.name, ex);
+    }
+  }
+
+  protected AbstractInsnNode processInstruction(
+      AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     switch (node.getOpcode()) {
       case Opcodes.RET:
       case Opcodes.RETURN:
@@ -179,7 +223,7 @@ protected AbstractInsnNode processInstruction(AbstractInsnNode node) {
       case Opcodes.ARETURN:
         {
           // stack [ret_value]
-          InsnList beforeReturnInsnList = getBeforeReturnInsnList(node);
+          InsnList beforeReturnInsnList = getBeforeReturnInsnList(node, frames);
           if (beforeReturnInsnList != null) {
             methodNode.instructions.insertBefore(node, beforeReturnInsnList);
           }
@@ -193,7 +237,8 @@ protected AbstractInsnNode processInstruction(AbstractInsnNode node) {
     return node;
   }
 
-  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+  protected InsnList getBeforeReturnInsnList(
+      AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     return null;
   }
 

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

@@ -56,6 +56,7 @@
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
@@ -71,6 +72,8 @@
 import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.TypeInsnNode;
 import org.objectweb.asm.tree.VarInsnNode;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -158,15 +161,18 @@ private InsnList wrapTryCatch(InsnList insnList) {
   }
 
   @Override
-  protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
+  protected InsnList getBeforeReturnInsnList(
+      AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     int size = 1;
     int storeOpCode = 0;
     int loadOpCode = 0;
     Type returnType = null;
     switch (node.getOpcode()) {
       case Opcodes.RET:
       case Opcodes.RETURN:
-        return wrapTryCatch(callMetric(metricProbe, node));
+        InsnList insnList = wrapTryCatch(callMetric(metricProbe, node));
+        insnList.insert(stackCleanupInsnList(node, 0, frames)); // void return: nothing to keep
+        return insnList;
       case Opcodes.LRETURN:
         storeOpCode = Opcodes.LSTORE;
         loadOpCode = Opcodes.LLOAD;
@@ -202,7 +208,10 @@ protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
         wrapTryCatch(
             callMetric(metricProbe, node, new ReturnContext(tmpIdx, loadOpCode, returnType)));
     // store return value from the stack to local before wrapped call
-    insnList.insert(new VarInsnNode(storeOpCode, tmpIdx));
+    InsnList prefixInsns = new InsnList();
+    prefixInsns.add(new VarInsnNode(storeOpCode, tmpIdx));
+    prefixInsns.add(stackCleanupInsnList(node, 1, frames)); // keep 1 value (the one we just stored)
+    insnList.insert(prefixInsns);
     // restore return value to the stack after wrapped call
     insnList.add(new VarInsnNode(loadOpCode, tmpIdx));
     return insnList;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/MetricProbesInstrumentationTest.java

@@ -31,10 +31,12 @@
 import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
+import java.io.File;
 import java.io.IOException;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
 import java.net.URISyntaxException;
+import java.net.URL;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -456,6 +458,33 @@ public void methodSyntheticDurationGaugeMetric() throws IOException, URISyntaxEx
     assertArrayEquals(new String[] {METRIC_PROBEID_TAG}, listener.lastTags);
   }
 
+  @Test
+  public void methodSyntheticDurationKotlinDist() {
+    final String CLASS_NAME = "CapturedSnapshot302";
+    String METRIC_NAME = "syn_dist";
+    MetricProbe metricProbe =
+        createMetricBuilder(METRIC_ID, METRIC_NAME, DISTRIBUTION)
+            .where(CLASS_NAME, "download", null)
+            .valueScript(new ValueScript(DSL.ref("@duration"), "@duration"))
+            .evaluateAt(MethodLocation.EXIT)
+            .build();
+    MetricForwarderListener listener = installMetricProbes(metricProbe);
+    URL resource = CapturedSnapshotTest.class.getResource("/" + CLASS_NAME + ".kt");
+    List<File> filesToDelete = new ArrayList<>();
+    try {
+      Class<?> testClass =
+          KotlinHelper.compileAndLoad(CLASS_NAME, resource.getFile(), filesToDelete);
+      Object companion = Reflect.onClass(testClass).get("Companion");
+      int result = Reflect.on(companion).call("main", "1").get();
+      assertEquals(1, result);
+      assertTrue(listener.doubleDistributions.containsKey(METRIC_NAME));
+      assertTrue(listener.doubleDistributions.get(METRIC_NAME).doubleValue() > 0);
+      assertArrayEquals(new String[] {METRIC_PROBEID_TAG}, listener.lastTags);
+    } finally {
+      filesToDelete.forEach(File::delete);
+    }
+  }
+
   @Test
   public void methodSyntheticDurationExceptionGaugeMetric() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot05";

dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/LogProbesIntegrationTest.java

@@ -356,7 +356,7 @@ void testLineProbe() throws Exception {
     LogProbe probe =
         LogProbe.builder()
             .probeId(LINE_PROBE_ID1)
-            .where("DebuggerTestApplication.java", 88)
+            .where("DebuggerTestApplication.java", 95)
             .captureSnapshot(true)
             .build();
     setCurrentConfiguration(createConfig(probe));
@@ -365,7 +365,7 @@ void testLineProbe() throws Exception {
     registerSnapshotListener(
         snapshot -> {
           assertEquals(LINE_PROBE_ID1.getId(), snapshot.getProbe().getId());
-          CapturedContext capturedContext = snapshot.getCaptures().getLines().get(88);
+          CapturedContext capturedContext = snapshot.getCaptures().getLines().get(95);
           assertFullMethodCaptureArgs(capturedContext);
           assertNull(capturedContext.getLocals());
           assertNull(capturedContext.getCapturedThrowable());

dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/MetricProbesIntegrationTest.java

@@ -210,7 +210,7 @@ private void doLineMetric(
         MetricProbe.builder()
             .probeId(PROBE_ID)
             // on line: System.out.println("fullMethod");
-            .where("DebuggerTestApplication.java", 88)
+            .where("DebuggerTestApplication.java", 95)
             .kind(kind)
             .metricName(metricName)
             .valueScript(script)