Skip to content

Commit 56d4919

Browse files
jpbempeljpbempel
authored
Add Exception probe custom instrumentation (#8230)
When there is only exception probes at a specific ode location we strip down the instrumentation to only capture values inside the catch (when an exception happens). this way, in normal execution no overhead is pay for an exception probe
1 parent afbea0b commit 56d4919

File tree

6 files changed

+138
-19
lines changed

6 files changed

+138
-19
lines changed

‎dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/CapturedContext.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,7 @@ public Status evaluate(
299299
MethodLocation methodLocation) {
300300
Status status =
301301
statusByProbeId.computeIfAbsent(encodedProbeId, key -> probeImplementation.createStatus());
302-
if (methodLocation == MethodLocation.EXIT) {
302+
if (methodLocation == MethodLocation.EXIT && startTimestamp > 0) {
303303
duration = System.nanoTime() - startTimestamp;
304304
addExtension(
305305
ValueReferences.DURATION_EXTENSION_NAME, duration / 1_000_000.0); // convert to ms

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -430,7 +430,6 @@ private byte[] writeClassFile(
430430
classNode.version = Opcodes.V1_8;
431431
}
432432
ClassWriter writer = new SafeClassWriter(loader);
433-
434433
log.debug("Generating bytecode for class: {}", Strings.getClassName(classFilePath));
435434
try {
436435
classNode.accept(writer);
@@ -695,6 +694,8 @@ private static void processCapturedContextLineProbes(
695694
private ProbeDefinition selectReferenceDefinition(
696695
List<ProbeDefinition> capturedContextProbes, ClassFileLines classFileLines) {
697696
boolean hasLogProbe = false;
697+
boolean hasOnlyExceptionProbe =
698+
capturedContextProbes.stream().allMatch(def -> def instanceof ExceptionProbe);
698699
MethodLocation evaluateAt = MethodLocation.EXIT;
699700
LogProbe.Capture capture = null;
700701
boolean captureSnapshot = false;
@@ -719,6 +720,10 @@ private ProbeDefinition selectReferenceDefinition(
719720
evaluateAt = definition.getEvaluateAt();
720721
}
721722
}
723+
if (hasOnlyExceptionProbe) {
724+
// only exception probes return the first one
725+
return capturedContextProbes.get(0);
726+
}
722727
if (hasLogProbe) {
723728
return LogProbe.builder()
724729
.probeId(probeId)

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

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ public InstrumentationResult.Status instrument() {
104104
instrumentMethodEnter();
105105
instrumentTryCatchHandlers();
106106
processInstructions();
107-
addFinallyHandler(returnHandlerLabel);
107+
addFinallyHandler(contextInitLabel, returnHandlerLabel);
108108
installFinallyBlocks();
109109
return InstrumentationResult.Status.INSTALLED;
110110
}
@@ -317,7 +317,11 @@ protected InsnList getReturnHandlerInsnList() {
317317
private InsnList commit() {
318318
InsnList insnList = new InsnList();
319319
// stack []
320-
getContext(insnList, entryContextVar);
320+
if (entryContextVar != -1) {
321+
getContext(insnList, entryContextVar);
322+
} else {
323+
getStatic(insnList, CAPTURED_CONTEXT_TYPE, "EMPTY_CAPTURING_CONTEXT");
324+
}
321325
// stack [capturedcontext]
322326
getContext(insnList, exitContextVar);
323327
// stack [capturedcontext, capturedcontext]
@@ -342,7 +346,7 @@ private InsnList commit() {
342346
return insnList;
343347
}
344348

345-
private void addFinallyHandler(LabelNode endLabel) {
349+
protected void addFinallyHandler(LabelNode startLabel, LabelNode endLabel) {
346350
// stack: [exception]
347351
if (methodNode.tryCatchBlocks == null) {
348352
methodNode.tryCatchBlocks = new ArrayList<>();
@@ -351,18 +355,28 @@ private void addFinallyHandler(LabelNode endLabel) {
351355
InsnList handler = new InsnList();
352356
handler.add(handlerLabel);
353357
// stack [exception]
354-
handler.add(new VarInsnNode(Opcodes.ALOAD, entryContextVar));
355-
// stack [exception, capturedcontext]
356-
LabelNode targetNode = new LabelNode();
357-
invokeVirtual(handler, CAPTURED_CONTEXT_TYPE, "isCapturing", BOOLEAN_TYPE);
358-
// stack [exception, boolean]
359-
handler.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
358+
LabelNode targetNode = null;
359+
if (entryContextVar != -1) {
360+
handler.add(new VarInsnNode(Opcodes.ALOAD, entryContextVar));
361+
// stack [exception, capturedcontext]
362+
targetNode = new LabelNode();
363+
invokeVirtual(handler, CAPTURED_CONTEXT_TYPE, "isCapturing", BOOLEAN_TYPE);
364+
// stack [exception, boolean]
365+
handler.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
366+
}
367+
if (exitContextVar == -1) {
368+
exitContextVar = newVar(CAPTURED_CONTEXT_TYPE);
369+
}
360370
// stack [exception]
361371
handler.add(collectCapturedContext(Snapshot.Kind.UNHANDLED_EXCEPTION, endLabel));
362372
// stack: [exception, capturedcontext]
363373
ldc(handler, Type.getObjectType(classNode.name));
364374
// stack [exception, capturedcontext, class]
365-
handler.add(new VarInsnNode(Opcodes.LLOAD, timestampStartVar));
375+
if (timestampStartVar != -1) {
376+
handler.add(new VarInsnNode(Opcodes.LLOAD, timestampStartVar));
377+
} else {
378+
ldc(handler, -1L);
379+
}
366380
// stack [exception, capturedcontext, class, long]
367381
getStatic(handler, METHOD_LOCATION_TYPE, "EXIT");
368382
// stack [exception, capturedcontext, class, long, methodlocation]
@@ -382,12 +396,14 @@ private void addFinallyHandler(LabelNode endLabel) {
382396
invokeStatic(handler, DEBUGGER_CONTEXT_TYPE, "disableInProbe", VOID_TYPE);
383397
// stack [exception]
384398
handler.add(commit());
385-
handler.add(targetNode);
399+
if (targetNode != null) {
400+
handler.add(targetNode);
401+
}
386402
// stack [exception]
387403
handler.add(new InsnNode(Opcodes.ATHROW));
388404
// stack: []
389405
methodNode.instructions.add(handler);
390-
finallyBlocks.add(new FinallyBlock(contextInitLabel, endLabel, handlerLabel));
406+
finallyBlocks.add(new FinallyBlock(startLabel, endLabel, handlerLabel));
391407
}
392408

393409
private void instrumentMethodEnter() {
@@ -921,7 +937,10 @@ private void collectLocalVariables(AbstractInsnNode location, InsnList insnList)
921937
}
922938
}
923939
}
924-
940+
if (applicableVars.isEmpty()) {
941+
// no applicable local variables - bail out
942+
return;
943+
}
925944
insnList.add(new InsnNode(Opcodes.DUP));
926945
// stack: [capturedcontext, capturedcontext]
927946
ldc(insnList, applicableVars.size());

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package com.datadog.debugger.instrumentation;
2+
3+
import com.datadog.debugger.probe.ProbeDefinition;
4+
import datadog.trace.bootstrap.debugger.Limits;
5+
import datadog.trace.bootstrap.debugger.ProbeId;
6+
import java.util.List;
7+
import org.objectweb.asm.tree.AbstractInsnNode;
8+
import org.objectweb.asm.tree.InsnList;
9+
10+
public class ExceptionInstrumentor extends CapturedContextInstrumentor {
11+
12+
public ExceptionInstrumentor(
13+
ProbeDefinition definition,
14+
MethodInfo methodInfo,
15+
List<DiagnosticMessage> diagnostics,
16+
List<ProbeId> probeIds) {
17+
super(definition, methodInfo, diagnostics, probeIds, true, false, Limits.DEFAULT);
18+
}
19+
20+
@Override
21+
public InstrumentationResult.Status instrument() {
22+
processInstructions(); // fill returnHandlerLabel
23+
addFinallyHandler(methodEnterLabel, returnHandlerLabel);
24+
installFinallyBlocks();
25+
return InstrumentationResult.Status.INSTALLED;
26+
}
27+
28+
@Override
29+
protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
30+
return null;
31+
}
32+
33+
@Override
34+
protected InsnList getReturnHandlerInsnList() {
35+
return new InsnList();
36+
}
37+
}

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ExceptionProbe.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77
import com.datadog.debugger.el.ProbeCondition;
88
import com.datadog.debugger.exception.ExceptionProbeManager;
99
import com.datadog.debugger.exception.Fingerprinter;
10+
import com.datadog.debugger.instrumentation.DiagnosticMessage;
11+
import com.datadog.debugger.instrumentation.ExceptionInstrumentor;
12+
import com.datadog.debugger.instrumentation.InstrumentationResult;
1013
import com.datadog.debugger.instrumentation.MethodInfo;
1114
import com.datadog.debugger.sink.Snapshot;
1215
import datadog.trace.bootstrap.debugger.CapturedContext;
@@ -48,6 +51,12 @@ public ExceptionProbe(
4851
this.chainedExceptionIdx = chainedExceptionIdx;
4952
}
5053

54+
@Override
55+
public InstrumentationResult.Status instrument(
56+
MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
57+
return new ExceptionInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
58+
}
59+
5160
@Override
5261
public boolean isLineProbe() {
5362
// Exception probe are always method probe even if there is a line number
@@ -62,7 +71,11 @@ public CapturedContext.Status createStatus() {
6271
@Override
6372
public void evaluate(
6473
CapturedContext context, CapturedContext.Status status, MethodLocation methodLocation) {
65-
if (!(status instanceof ExceptionProbeStatus)) {
74+
ExceptionProbeStatus exceptionStatus;
75+
if (status instanceof ExceptionProbeStatus) {
76+
exceptionStatus = (ExceptionProbeStatus) status;
77+
exceptionStatus.setCapture(false);
78+
} else {
6679
throw new IllegalStateException("Invalid status: " + status.getClass());
6780
}
6881
if (methodLocation != MethodLocation.EXIT) {
@@ -82,7 +95,6 @@ public void evaluate(
8295
if (exceptionProbeManager.shouldCaptureException(fingerprint)) {
8396
LOGGER.debug("Capturing exception matching fingerprint: {}", fingerprint);
8497
// capture only on uncaught exception matching the fingerprint
85-
ExceptionProbeStatus exceptionStatus = (ExceptionProbeStatus) status;
8698
ExceptionProbeManager.ThrowableState state =
8799
exceptionProbeManager.getStateByThrowable(innerMostThrowable);
88100
if (state != null) {
@@ -148,7 +160,7 @@ public void buildLocation(MethodInfo methodInfo) {
148160
}
149161

150162
public static class ExceptionProbeStatus extends LogStatus {
151-
private boolean capture;
163+
private boolean capture = true; // default to true for status entry when mixed with log probe
152164

153165
public ExceptionProbeStatus(ProbeImplementation probeImplementation) {
154166
super(probeImplementation);

‎dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/ExceptionProbeInstrumentationTest.java‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import com.datadog.debugger.agent.JsonSnapshotSerializer;
2323
import com.datadog.debugger.agent.MockSampler;
2424
import com.datadog.debugger.probe.ExceptionProbe;
25+
import com.datadog.debugger.probe.LogProbe;
26+
import com.datadog.debugger.probe.ProbeDefinition;
2527
import com.datadog.debugger.sink.DebuggerSink;
2628
import com.datadog.debugger.sink.ProbeStatusSink;
2729
import com.datadog.debugger.sink.Snapshot;
@@ -33,14 +35,19 @@
3335
import datadog.trace.api.interceptor.MutableSpan;
3436
import datadog.trace.bootstrap.debugger.DebuggerContext;
3537
import datadog.trace.bootstrap.debugger.DebuggerContext.ClassNameFilter;
38+
import datadog.trace.bootstrap.debugger.ProbeId;
3639
import datadog.trace.bootstrap.debugger.ProbeLocation;
3740
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
3841
import datadog.trace.core.CoreTracer;
42+
import java.io.IOException;
3943
import java.lang.instrument.ClassFileTransformer;
4044
import java.lang.instrument.Instrumentation;
45+
import java.net.URISyntaxException;
4146
import java.time.Clock;
4247
import java.time.Duration;
4348
import java.time.Instant;
49+
import java.util.Arrays;
50+
import java.util.Collection;
4451
import java.util.Map;
4552
import java.util.Set;
4653
import java.util.stream.Collectors;
@@ -54,6 +61,8 @@
5461
import org.junit.jupiter.api.condition.DisabledIf;
5562

5663
public class ExceptionProbeInstrumentationTest {
64+
protected static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f5", 0);
65+
5766
private final Instrumentation instr = ByteBuddyAgent.install();
5867
private final TestTraceInterceptor traceInterceptor = new TestTraceInterceptor();
5968
private ClassFileTransformer currentTransformer;
@@ -84,6 +93,7 @@ public void before() {
8493
ProbeRateLimiter.setGlobalSnapshotRate(1000);
8594
// to activate the call to DebuggerContext.handleException
8695
setFieldInConfig(Config.get(), "debuggerExceptionEnabled", true);
96+
setFieldInConfig(Config.get(), "debuggerClassFileDumpEnabled", true);
8797
}
8898

8999
@AfterEach
@@ -252,6 +262,34 @@ public void captureOncePerHour() throws Exception {
252262
assertEquals(1, listener.snapshots.size());
253263
}
254264

265+
@Test
266+
public void mixedWithLogProbes() throws IOException, URISyntaxException {
267+
final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
268+
Config config = createConfig();
269+
ExceptionProbeManager exceptionProbeManager = new ExceptionProbeManager(classNameFiltering);
270+
LogProbe logProbe =
271+
LogProbe.builder().probeId(PROBE_ID).where(CLASS_NAME, "processWithException").build();
272+
Collection<ProbeDefinition> definitions = Arrays.asList(logProbe);
273+
TestSnapshotListener listener =
274+
setupExceptionDebugging(config, exceptionProbeManager, classNameFiltering, definitions);
275+
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
276+
String fingerprint =
277+
callMethodThrowingRuntimeException(testClass); // instrument exception stacktrace
278+
assertWithTimeout(
279+
() -> exceptionProbeManager.isAlreadyInstrumented(fingerprint), Duration.ofSeconds(30));
280+
assertEquals(2, exceptionProbeManager.getProbes().size());
281+
callMethodThrowingRuntimeException(testClass); // generate snapshots
282+
Map<String, Set<String>> probeIdsByMethodName =
283+
extractProbeIdsByMethodName(exceptionProbeManager);
284+
assertEquals(3, listener.snapshots.size()); // 2 log snapshots + 1 exception snapshot
285+
Snapshot snapshot0 = listener.snapshots.get(0);
286+
assertEquals(PROBE_ID.getId(), snapshot0.getProbe().getId());
287+
Snapshot snapshot1 = listener.snapshots.get(1);
288+
assertEquals(PROBE_ID.getId(), snapshot1.getProbe().getId());
289+
Snapshot snapshot2 = listener.snapshots.get(2);
290+
assertProbeId(probeIdsByMethodName, "processWithException", snapshot2.getProbe().getId());
291+
}
292+
255293
private static void assertExceptionMsg(String expectedMsg, Snapshot snapshot) {
256294
assertEquals(
257295
expectedMsg, snapshot.getCaptures().getReturn().getCapturedThrowable().getMessage());
@@ -314,6 +352,14 @@ private TestSnapshotListener setupExceptionDebugging(
314352
Config config,
315353
ExceptionProbeManager exceptionProbeManager,
316354
ClassNameFilter classNameFiltering) {
355+
return setupExceptionDebugging(config, exceptionProbeManager, classNameFiltering, null);
356+
}
357+
358+
private TestSnapshotListener setupExceptionDebugging(
359+
Config config,
360+
ExceptionProbeManager exceptionProbeManager,
361+
ClassNameFilter classNameFiltering,
362+
Collection<ProbeDefinition> definitions) {
317363
ProbeStatusSink probeStatusSink = mock(ProbeStatusSink.class);
318364
ConfigurationUpdater configurationUpdater =
319365
new ConfigurationUpdater(
@@ -330,7 +376,7 @@ private TestSnapshotListener setupExceptionDebugging(
330376
new DefaultExceptionDebugger(
331377
exceptionProbeManager, configurationUpdater, classNameFiltering, 100);
332378
DebuggerContext.initExceptionDebugger(exceptionDebugger);
333-
configurationUpdater.accept(REMOTE_CONFIG, null);
379+
configurationUpdater.accept(REMOTE_CONFIG, definitions);
334380
return listener;
335381
}
336382

0 commit comments

Comments
 (0)