Fix suspend Kotlin method instrumentation

suspend method in Kotlin generates a special bytecode for the method
with a state machine and could return different objects.
for each return branch we need to have a specific return handler
to avoid having mismatch stack maps

Author: jpbempel

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

@@ -150,12 +150,15 @@ private AbstractInsnNode findFirstInsnForConstructor(AbstractInsnNode first) {
 
   protected void processInstructions() {
     AbstractInsnNode node = methodNode.instructions.getFirst();
-    while (node != null && !node.equals(returnHandlerLabel)) {
+    LabelNode sentinelNode = new LabelNode();
+    methodNode.instructions.add(sentinelNode);
+    while (node != null && !node.equals(sentinelNode)) {
       if (node.getType() != AbstractInsnNode.LINE) {
         node = processInstruction(node);
       }
       node = node.getNext();
     }
+    methodNode.instructions.remove(sentinelNode);
     if (returnHandlerLabel == null) {
       // if no return found, fallback to use the last instruction as last resort
       returnHandlerLabel = new LabelNode();
@@ -197,9 +200,8 @@ protected LabelNode getReturnHandler(AbstractInsnNode exitNode) {
     if (exitNode.getNext() != null || exitNode.getPrevious() != null) {
       throw new IllegalArgumentException("exitNode is not removed from original instruction list");
     }
-    if (returnHandlerLabel != null) {
-      return returnHandlerLabel;
-    }
+    // Create the returnHandlerLabel every time because the stack state could be different
+    // for each return (suspend method in Kotlin)
     returnHandlerLabel = new LabelNode();
     methodNode.instructions.add(returnHandlerLabel);
     // stack top is return value (if any)

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

@@ -589,17 +589,27 @@ public void suspendKotlin() {
   public void suspendMethodKotlin() {
     final String CLASS_NAME = "CapturedSnapshot302";
     TestSnapshotListener listener =
-        installProbes(createProbe(PROBE_ID, CLASS_NAME , "download", null));
+        installProbes(createProbe(PROBE_ID, CLASS_NAME, "download", null));
     URL resource = CapturedSnapshotTest.class.getResource("/" + CLASS_NAME + ".kt");
     assertNotNull(resource);
     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", "").get();
-      assertEquals(0, result);
-      Snapshot snapshot = assertOneSnapshot(listener);
+      int result = Reflect.on(companion).call("main", "1").get();
+      assertEquals(1, result);
+      // 2 snapshots are expected because the method is executed twice one for each state
+      // before the delay, after the delay
+      List<Snapshot> snapshots = assertSnapshots(listener, 2);
+      Snapshot snapshot0 = snapshots.get(0);
+      assertCaptureReturnValue(
+          snapshot0.getCaptures().getReturn(),
+          "kotlin.coroutines.intrinsics.CoroutineSingletons",
+          "COROUTINE_SUSPENDED");
+      Snapshot snapshot1 = snapshots.get(1);
+      assertCaptureReturnValue(
+          snapshot1.getCaptures().getReturn(), String.class.getTypeName(), "1");
     } finally {
       filesToDelete.forEach(File::delete);
     }