Fix suspend Kotlin methods instrumentation (#8080)

suspend methods 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
add config for local var hoisting `DD_DYNAMIC_INSTRUMENTATION_HOIST_LOCALVARS_ENABLED`
enable only hoisting for java language

Author: jpbempel

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

@@ -396,7 +396,11 @@ private void instrumentMethodEnter() {
     if (methodNode.tryCatchBlocks.size() > 0) {
       throwableListVar = declareThrowableList(insnList);
     }
-    unscopedLocalVars = initAndHoistLocalVars(insnList);
+    unscopedLocalVars = Collections.emptyList();
+    if (Config.get().isDebuggerHoistLocalVarsEnabled() && language == JvmLanguage.JAVA) {
+      // for now, only hoist local vars for Java
+      unscopedLocalVars = initAndHoistLocalVars(insnList);
+    }
     insnList.add(contextInitLabel);
     if (definition instanceof SpanDecorationProbe
         && definition.getEvaluateAt() == MethodLocation.EXIT) {

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

@@ -46,6 +46,7 @@ public abstract class Instrumentor {
   protected int localVarBaseOffset;
   protected int argOffset;
   protected final LocalVariableNode[] localVarsBySlotArray;
+  protected final JvmLanguage language;
   protected LabelNode returnHandlerLabel;
   protected final List<CapturedContextInstrumentor.FinallyBlock> finallyBlocks = new ArrayList<>();
 
@@ -69,6 +70,7 @@ public Instrumentor(
       argOffset += t.getSize();
     }
     localVarsBySlotArray = extractLocalVariables(argTypes);
+    this.language = JvmLanguage.of(classNode);
   }
 
   public abstract InstrumentationResult.Status instrument();
@@ -150,12 +152,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 +202,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)
@@ -292,4 +296,31 @@ public FinallyBlock(LabelNode startLabel, LabelNode endLabel, LabelNode handlerL
       this.handlerLabel = handlerLabel;
     }
   }
+
+  protected enum JvmLanguage {
+    JAVA,
+    KOTLIN,
+    SCALA,
+    GROOVY,
+    UNKNOWN;
+
+    public static JvmLanguage of(ClassNode classNode) {
+      if (classNode.sourceFile == null) {
+        return UNKNOWN;
+      }
+      if (classNode.sourceFile.endsWith(".java")) {
+        return JAVA;
+      }
+      if (classNode.sourceFile.endsWith(".kt")) {
+        return KOTLIN;
+      }
+      if (classNode.sourceFile.endsWith(".scala")) {
+        return SCALA;
+      }
+      if (classNode.sourceFile.endsWith(".groovy")) {
+        return GROOVY;
+      }
+      return UNKNOWN;
+    }
+  }
 }

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

@@ -582,6 +582,39 @@ public void suspendKotlin() {
     }
   }
 
+  @Test
+  @DisabledIf(
+      value = "datadog.trace.api.Platform#isJ9",
+      disabledReason = "Issue with J9 when compiling Kotlin code")
+  public void suspendMethodKotlin() {
+    final String CLASS_NAME = "CapturedSnapshot302";
+    TestSnapshotListener listener =
+        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", "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);
+    }
+  }
+
   @Test
   @DisabledIf(
       value = "datadog.trace.api.Platform#isJ9",

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -181,6 +181,7 @@ public final class ConfigDefaults {
   static final boolean DEFAULT_DEBUGGER_VERIFY_BYTECODE = true;
   static final boolean DEFAULT_DEBUGGER_INSTRUMENT_THE_WORLD = false;
   static final int DEFAULT_DEBUGGER_CAPTURE_TIMEOUT = 100; // milliseconds
+  static final boolean DEFAULT_DEBUGGER_HOIST_LOCALVARS_ENABLED = true;
   static final boolean DEFAULT_DEBUGGER_SYMBOL_ENABLED = true;
   static final boolean DEFAULT_DEBUGGER_SYMBOL_FORCE_UPLOAD = false;
   static final int DEFAULT_DEBUGGER_SYMBOL_FLUSH_THRESHOLD = 100; // nb of classes

dd-trace-api/src/main/java/datadog/trace/api/config/DebuggerConfig.java

@@ -28,6 +28,8 @@ public final class DebuggerConfig {
   public static final String DEBUGGER_REDACTION_EXCLUDED_IDENTIFIERS =
       "dynamic.instrumentation.redaction.excluded.identifiers";
   public static final String DEBUGGER_REDACTED_TYPES = "dynamic.instrumentation.redacted.types";
+  public static final String DEBUGGER_HOIST_LOCALVARS_ENABLED =
+      "dynamic.instrumentation.hoist.localvars.enabled";
   public static final String DEBUGGER_SYMBOL_ENABLED = "symbol.database.upload.enabled";
   public static final String DEBUGGER_SYMBOL_FORCE_UPLOAD = "internal.force.symbol.database.upload";
   public static final String DEBUGGER_SYMBOL_INCLUDES = "symbol.database.includes";

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -396,6 +396,7 @@ public static String getHostName() {
   private final String debuggerRedactedIdentifiers;
   private final Set<String> debuggerRedactionExcludedIdentifiers;
   private final String debuggerRedactedTypes;
+  private final boolean debuggerHoistLocalVarsEnabled;
   private final boolean debuggerSymbolEnabled;
   private final boolean debuggerSymbolForceUpload;
   private final String debuggerSymbolIncludes;
@@ -1538,6 +1539,9 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     debuggerRedactionExcludedIdentifiers =
         tryMakeImmutableSet(configProvider.getList(DEBUGGER_REDACTION_EXCLUDED_IDENTIFIERS));
     debuggerRedactedTypes = configProvider.getString(DEBUGGER_REDACTED_TYPES, null);
+    debuggerHoistLocalVarsEnabled =
+        configProvider.getBoolean(
+            DEBUGGER_HOIST_LOCALVARS_ENABLED, DEFAULT_DEBUGGER_HOIST_LOCALVARS_ENABLED);
     debuggerSymbolEnabled =
         configProvider.getBoolean(DEBUGGER_SYMBOL_ENABLED, DEFAULT_DEBUGGER_SYMBOL_ENABLED);
     debuggerSymbolForceUpload =
@@ -3063,6 +3067,10 @@ public String getDebuggerRedactedTypes() {
     return debuggerRedactedTypes;
   }
 
+  public boolean isDebuggerHoistLocalVarsEnabled() {
+    return debuggerHoistLocalVarsEnabled;
+  }
+
   public boolean isAwsPropagationEnabled() {
     return awsPropagationEnabled;
   }