From baedf8dd5796e6b06e509eed7e90417e553cc955 Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bempel Date: Fri, 13 Dec 2024 11:32:27 +0100 Subject: [PATCH] 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 --- .../CapturedContextInstrumentor.java | 6 ++- .../instrumentation/Instrumentor.java | 39 +++++++++++++++++-- .../debugger/agent/CapturedSnapshotTest.java | 33 ++++++++++++++++ .../datadog/trace/api/ConfigDefaults.java | 1 + .../trace/api/config/DebuggerConfig.java | 2 + .../main/java/datadog/trace/api/Config.java | 8 ++++ 6 files changed, 84 insertions(+), 5 deletions(-) diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java index 32b8cce5804..48808cd0b2b 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java +++ b/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) { diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumentor.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumentor.java index d16e2fda812..358c820b049 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumentor.java +++ b/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 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; + } + } } diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java index e8dc10ab2b7..ed712ca000d 100644 --- a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java +++ b/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 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 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", diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index a3f2915dcb2..4b3df14e981 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/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 diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/DebuggerConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/DebuggerConfig.java index d0c26d01ed1..aa0619121e9 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/DebuggerConfig.java +++ b/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"; diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index e24fe3a2710..f27c6cf0272 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/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 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; }