Skip to content

Commit

Permalink
Add Exception probe custom instrumentation
Browse files Browse the repository at this point in the history
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
  • Loading branch information
jpbempel committed Jan 16, 2025
1 parent 8b63e5a commit 7ab6064
Show file tree
Hide file tree
Showing 6 changed files with 138 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -299,7 +299,7 @@ public Status evaluate(
MethodLocation methodLocation) {
Status status =
statusByProbeId.computeIfAbsent(encodedProbeId, key -> probeImplementation.createStatus());
if (methodLocation == MethodLocation.EXIT) {
if (methodLocation == MethodLocation.EXIT && startTimestamp > 0) {
duration = System.nanoTime() - startTimestamp;
addExtension(
ValueReferences.DURATION_EXTENSION_NAME, duration / 1_000_000.0); // convert to ms
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -428,7 +428,6 @@ private byte[] writeClassFile(
classNode.version = Opcodes.V1_8;
}
ClassWriter writer = new SafeClassWriter(loader);

log.debug("Generating bytecode for class: {}", Strings.getClassName(classFilePath));
try {
classNode.accept(writer);
Expand Down Expand Up @@ -691,6 +690,8 @@ private static void processCapturedContextLineProbes(
private ProbeDefinition selectReferenceDefinition(
List<ProbeDefinition> capturedContextProbes, ClassFileLines classFileLines) {
boolean hasLogProbe = false;
boolean hasOnlyExceptionProbe =
capturedContextProbes.stream().allMatch(def -> def instanceof ExceptionProbe);
MethodLocation evaluateAt = MethodLocation.EXIT;
LogProbe.Capture capture = null;
boolean captureSnapshot = false;
Expand All @@ -715,6 +716,10 @@ private ProbeDefinition selectReferenceDefinition(
evaluateAt = definition.getEvaluateAt();
}
}
if (hasOnlyExceptionProbe) {
// only exception probes return the first one
return capturedContextProbes.get(0);
}
if (hasLogProbe) {
return LogProbe.builder()
.probeId(probeId)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ public InstrumentationResult.Status instrument() {
instrumentMethodEnter();
instrumentTryCatchHandlers();
processInstructions();
addFinallyHandler(returnHandlerLabel);
addFinallyHandler(contextInitLabel, returnHandlerLabel);
installFinallyBlocks();
return InstrumentationResult.Status.INSTALLED;
}
Expand Down Expand Up @@ -317,7 +317,11 @@ protected InsnList getReturnHandlerInsnList() {
private InsnList commit() {
InsnList insnList = new InsnList();
// stack []
getContext(insnList, entryContextVar);
if (entryContextVar != -1) {
getContext(insnList, entryContextVar);
} else {
getStatic(insnList, CAPTURED_CONTEXT_TYPE, "EMPTY_CAPTURING_CONTEXT");
}
// stack [capturedcontext]
getContext(insnList, exitContextVar);
// stack [capturedcontext, capturedcontext]
Expand All @@ -342,7 +346,7 @@ private InsnList commit() {
return insnList;
}

private void addFinallyHandler(LabelNode endLabel) {
protected void addFinallyHandler(LabelNode startLabel, LabelNode endLabel) {
// stack: [exception]
if (methodNode.tryCatchBlocks == null) {
methodNode.tryCatchBlocks = new ArrayList<>();
Expand All @@ -351,18 +355,28 @@ private void addFinallyHandler(LabelNode endLabel) {
InsnList handler = new InsnList();
handler.add(handlerLabel);
// stack [exception]
handler.add(new VarInsnNode(Opcodes.ALOAD, entryContextVar));
// stack [exception, capturedcontext]
LabelNode targetNode = new LabelNode();
invokeVirtual(handler, CAPTURED_CONTEXT_TYPE, "isCapturing", BOOLEAN_TYPE);
// stack [exception, boolean]
handler.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
LabelNode targetNode = null;
if (entryContextVar != -1) {
handler.add(new VarInsnNode(Opcodes.ALOAD, entryContextVar));
// stack [exception, capturedcontext]
targetNode = new LabelNode();
invokeVirtual(handler, CAPTURED_CONTEXT_TYPE, "isCapturing", BOOLEAN_TYPE);
// stack [exception, boolean]
handler.add(new JumpInsnNode(Opcodes.IFEQ, targetNode));
}
if (exitContextVar == -1) {
exitContextVar = newVar(CAPTURED_CONTEXT_TYPE);
}
// stack [exception]
handler.add(collectCapturedContext(Snapshot.Kind.UNHANDLED_EXCEPTION, endLabel));
// stack: [exception, capturedcontext]
ldc(handler, Type.getObjectType(classNode.name));
// stack [exception, capturedcontext, class]
handler.add(new VarInsnNode(Opcodes.LLOAD, timestampStartVar));
if (timestampStartVar != -1) {
handler.add(new VarInsnNode(Opcodes.LLOAD, timestampStartVar));
} else {
ldc(handler, -1L);
}
// stack [exception, capturedcontext, class, long]
getStatic(handler, METHOD_LOCATION_TYPE, "EXIT");
// stack [exception, capturedcontext, class, long, methodlocation]
Expand All @@ -382,12 +396,14 @@ private void addFinallyHandler(LabelNode endLabel) {
invokeStatic(handler, DEBUGGER_CONTEXT_TYPE, "disableInProbe", VOID_TYPE);
// stack [exception]
handler.add(commit());
handler.add(targetNode);
if (targetNode != null) {
handler.add(targetNode);
}
// stack [exception]
handler.add(new InsnNode(Opcodes.ATHROW));
// stack: []
methodNode.instructions.add(handler);
finallyBlocks.add(new FinallyBlock(contextInitLabel, endLabel, handlerLabel));
finallyBlocks.add(new FinallyBlock(startLabel, endLabel, handlerLabel));
}

private void instrumentMethodEnter() {
Expand Down Expand Up @@ -921,7 +937,10 @@ private void collectLocalVariables(AbstractInsnNode location, InsnList insnList)
}
}
}

if (applicableVars.isEmpty()) {
// no applicable local variables - bail out
return;
}
insnList.add(new InsnNode(Opcodes.DUP));
// stack: [capturedcontext, capturedcontext]
ldc(insnList, applicableVars.size());
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
package com.datadog.debugger.instrumentation;

import com.datadog.debugger.probe.ProbeDefinition;
import datadog.trace.bootstrap.debugger.Limits;
import datadog.trace.bootstrap.debugger.ProbeId;
import java.util.List;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.InsnList;

public class ExceptionInstrumentor extends CapturedContextInstrumentor {

public ExceptionInstrumentor(
ProbeDefinition definition,
MethodInfo methodInfo,
List<DiagnosticMessage> diagnostics,
List<ProbeId> probeIds) {
super(definition, methodInfo, diagnostics, probeIds, true, false, Limits.DEFAULT);
}

@Override
public InstrumentationResult.Status instrument() {
processInstructions(); // fill returnHandlerLabel
addFinallyHandler(methodEnterLabel, returnHandlerLabel);
installFinallyBlocks();
return InstrumentationResult.Status.INSTALLED;
}

@Override
protected InsnList getBeforeReturnInsnList(AbstractInsnNode node) {
return null;
}

@Override
protected InsnList getReturnHandlerInsnList() {
return new InsnList();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,10 @@
import com.datadog.debugger.el.ProbeCondition;
import com.datadog.debugger.exception.ExceptionProbeManager;
import com.datadog.debugger.exception.Fingerprinter;
import com.datadog.debugger.instrumentation.DiagnosticMessage;
import com.datadog.debugger.instrumentation.ExceptionInstrumentor;
import com.datadog.debugger.instrumentation.InstrumentationResult;
import com.datadog.debugger.instrumentation.MethodInfo;
import com.datadog.debugger.sink.Snapshot;
import datadog.trace.bootstrap.debugger.CapturedContext;
import datadog.trace.bootstrap.debugger.MethodLocation;
Expand Down Expand Up @@ -48,6 +51,12 @@ public ExceptionProbe(
this.chainedExceptionIdx = chainedExceptionIdx;
}

@Override
public InstrumentationResult.Status instrument(
MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
return new ExceptionInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
}

@Override
public boolean isLineProbe() {
// Exception probe are always method probe even if there is a line number
Expand All @@ -62,7 +71,11 @@ public CapturedContext.Status createStatus() {
@Override
public void evaluate(
CapturedContext context, CapturedContext.Status status, MethodLocation methodLocation) {
if (!(status instanceof ExceptionProbeStatus)) {
ExceptionProbeStatus exceptionStatus;
if (status instanceof ExceptionProbeStatus) {
exceptionStatus = (ExceptionProbeStatus) status;
exceptionStatus.setCapture(false);
} else {
throw new IllegalStateException("Invalid status: " + status.getClass());
}
if (methodLocation != MethodLocation.EXIT) {
Expand All @@ -82,7 +95,6 @@ public void evaluate(
if (exceptionProbeManager.shouldCaptureException(fingerprint)) {
LOGGER.debug("Capturing exception matching fingerprint: {}", fingerprint);
// capture only on uncaught exception matching the fingerprint
ExceptionProbeStatus exceptionStatus = (ExceptionProbeStatus) status;
ExceptionProbeManager.ThrowableState state =
exceptionProbeManager.getStateByThrowable(innerMostThrowable);
if (state != null) {
Expand Down Expand Up @@ -148,7 +160,7 @@ public void buildLocation(InstrumentationResult result) {
}

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

public ExceptionProbeStatus(ProbeImplementation probeImplementation) {
super(probeImplementation);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
import com.datadog.debugger.agent.JsonSnapshotSerializer;
import com.datadog.debugger.agent.MockSampler;
import com.datadog.debugger.probe.ExceptionProbe;
import com.datadog.debugger.probe.LogProbe;
import com.datadog.debugger.probe.ProbeDefinition;
import com.datadog.debugger.sink.DebuggerSink;
import com.datadog.debugger.sink.ProbeStatusSink;
import com.datadog.debugger.sink.Snapshot;
Expand All @@ -33,14 +35,19 @@
import datadog.trace.api.interceptor.MutableSpan;
import datadog.trace.bootstrap.debugger.DebuggerContext;
import datadog.trace.bootstrap.debugger.DebuggerContext.ClassNameFilter;
import datadog.trace.bootstrap.debugger.ProbeId;
import datadog.trace.bootstrap.debugger.ProbeLocation;
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
import datadog.trace.core.CoreTracer;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.net.URISyntaxException;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
Expand All @@ -54,6 +61,8 @@
import org.junit.jupiter.api.condition.DisabledIf;

public class ExceptionProbeInstrumentationTest {
protected static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f5", 0);

private final Instrumentation instr = ByteBuddyAgent.install();
private final TestTraceInterceptor traceInterceptor = new TestTraceInterceptor();
private ClassFileTransformer currentTransformer;
Expand Down Expand Up @@ -84,6 +93,7 @@ public void before() {
ProbeRateLimiter.setGlobalSnapshotRate(1000);
// to activate the call to DebuggerContext.handleException
setFieldInConfig(Config.get(), "debuggerExceptionEnabled", true);
setFieldInConfig(Config.get(), "debuggerClassFileDumpEnabled", true);
}

@AfterEach
Expand Down Expand Up @@ -252,6 +262,34 @@ public void captureOncePerHour() throws Exception {
assertEquals(1, listener.snapshots.size());
}

@Test
public void mixedWithLogProbes() throws IOException, URISyntaxException {
final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
Config config = createConfig();
ExceptionProbeManager exceptionProbeManager = new ExceptionProbeManager(classNameFiltering);
LogProbe logProbe =
LogProbe.builder().probeId(PROBE_ID).where(CLASS_NAME, "processWithException").build();
Collection<ProbeDefinition> definitions = Arrays.asList(logProbe);
TestSnapshotListener listener =
setupExceptionDebugging(config, exceptionProbeManager, classNameFiltering, definitions);
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
String fingerprint =
callMethodThrowingRuntimeException(testClass); // instrument exception stacktrace
assertWithTimeout(
() -> exceptionProbeManager.isAlreadyInstrumented(fingerprint), Duration.ofSeconds(30));
assertEquals(2, exceptionProbeManager.getProbes().size());
callMethodThrowingRuntimeException(testClass); // generate snapshots
Map<String, Set<String>> probeIdsByMethodName =
extractProbeIdsByMethodName(exceptionProbeManager);
assertEquals(3, listener.snapshots.size()); // 2 log snapshots + 1 exception snapshot
Snapshot snapshot0 = listener.snapshots.get(0);
assertEquals(PROBE_ID.getId(), snapshot0.getProbe().getId());
Snapshot snapshot1 = listener.snapshots.get(1);
assertEquals(PROBE_ID.getId(), snapshot1.getProbe().getId());
Snapshot snapshot2 = listener.snapshots.get(2);
assertProbeId(probeIdsByMethodName, "processWithException", snapshot2.getProbe().getId());
}

private static void assertExceptionMsg(String expectedMsg, Snapshot snapshot) {
assertEquals(
expectedMsg, snapshot.getCaptures().getReturn().getCapturedThrowable().getMessage());
Expand Down Expand Up @@ -314,6 +352,14 @@ private TestSnapshotListener setupExceptionDebugging(
Config config,
ExceptionProbeManager exceptionProbeManager,
ClassNameFilter classNameFiltering) {
return setupExceptionDebugging(config, exceptionProbeManager, classNameFiltering, null);
}

private TestSnapshotListener setupExceptionDebugging(
Config config,
ExceptionProbeManager exceptionProbeManager,
ClassNameFilter classNameFiltering,
Collection<ProbeDefinition> definitions) {
ProbeStatusSink probeStatusSink = mock(ProbeStatusSink.class);
ConfigurationUpdater configurationUpdater =
new ConfigurationUpdater(
Expand All @@ -330,7 +376,7 @@ private TestSnapshotListener setupExceptionDebugging(
new DefaultExceptionDebugger(
exceptionProbeManager, configurationUpdater, classNameFiltering, 100);
DebuggerContext.initExceptionDebugger(exceptionDebugger);
configurationUpdater.accept(REMOTE_CONFIG, null);
configurationUpdater.accept(REMOTE_CONFIG, definitions);
return listener;
}

Expand Down

0 comments on commit 7ab6064

Please sign in to comment.