Skip to content

Fix instrumenting metric probes for Kotlin #11890

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gh-worker-dd-mergequeue-cf854d merged 2 commits into from
Jul 9, 2026
+103 −12
Merged

Fix instrumenting metric probes for Kotlin #11890

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b063fca
Fix instrumenting metric probes for Kotlin
jpbempel Jul 8, 2026
61454ee
fix leftover
jpbempel Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion ...ugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumenter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -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));
Expand Down
6 changes: 5 additions & 1 deletion ...nt-debugger/src/main/java/com/datadog/debugger/instrumentation/ExceptionInstrumenter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand All @@ -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;
}

Expand Down
53 changes: 49 additions & 4 deletions ...agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumenter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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";

Expand Down Expand Up @@ -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();
}
Expand All @@ -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);
}

@ojung ojung Jul 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this produce a half-finished map? Would that be an issue?

@jpbempel jpbempel Jul 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No because the specific exception is thrown only by ASM's Analyzer.analyze method so when the exception is thrown, the map is untouchly empty

}

protected AbstractInsnNode processInstruction(
AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
switch (node.getOpcode()) {
case Opcodes.RET:
case Opcodes.RETURN:
Expand All @@ -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);
}
Expand All @@ -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;
}

Expand Down
15 changes: 12 additions & 3 deletions ...agent-debugger/src/main/java/com/datadog/debugger/instrumentation/MetricInstrumenter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -158,15 +161,18 @@ private InsnList wrapTryCatch(InsnList insnList) {
}

@Override
protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
protected InsnList getBeforeReturnInsnList(

@ojung ojung Jul 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we sure this is limited to metric probes?

@jpbempel jpbempel Jul 9, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Until contrary proven, yes because this is specific to the code inserted for metric probe instrumentation.
Though, it is not excluded to manifest in other instrumentation cases.
This is why we will add exploration tests for Kotlin code

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;
Expand Down Expand Up @@ -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;
Expand Down
29 changes: 29 additions & 0 deletions ...nt-debugger/src/test/java/com/datadog/debugger/agent/MetricProbesInstrumentationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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";
Expand Down
4 changes: 2 additions & 2 deletions .../debugger-integration-tests/src/test/java/datadog/smoketest/LogProbesIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion ...bugger-integration-tests/src/test/java/datadog/smoketest/MetricProbesIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading