Skip to content
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.

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

Add Exception probe custom instrumentation #8230

Merged
merged 1 commit into from
Jan 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
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 @@ -430,7 +430,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 @@ -695,6 +694,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 @@ -719,6 +720,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,6 +7,9 @@
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;
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);
Copy link
Contributor

Choose a reason for hiding this comment

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

why set false here?

Copy link
Member Author

Choose a reason for hiding this comment

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

because by default the status is capture true (for the entry part) then when evaluating it is turned to false until it reached the end of evaluation and this is the exception that we need to report

} 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(MethodInfo methodInfo) {
}

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
Loading