Skip to content

Commit

Permalink
Fix mixing log/span decoration probes (#7246)
Browse files Browse the repository at this point in the history 
On same location, if we have different probes like log or span
decoration probe with different characteristics (capture, snapshot,
evaluateAt,...), instrumentation needs the correct definition to apply
 the shared CapturedContextInstrumentor.
We are creating a synthetic definition for applying instrumentation
and merging all characteristics required for instrumentation.
  • Loading branch information
@jpbempel
jpbempel committed Jul 9, 2024
1 parent 5ab86d2 commit 3de43de
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 24 deletions.
76 changes: 62 additions & 14 deletions ...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.datadog.debugger.util.ExceptionHelper;
import datadog.trace.agent.tooling.AgentStrategies;
import datadog.trace.api.Config;
import datadog.trace.bootstrap.debugger.MethodLocation;
import datadog.trace.bootstrap.debugger.ProbeId;
import datadog.trace.bootstrap.debugger.ProbeImplementation;
import datadog.trace.util.Strings;
Expand Down Expand Up @@ -515,7 +516,8 @@ private InstrumentationResult applyInstrumentation(
InstrumentationResult.Status status = preCheckInstrumentation(diagnostics, methodInfo);
if (status != InstrumentationResult.Status.ERROR) {
try {
List<ToInstrumentInfo> toInstruments = filterAndSortDefinitions(definitions);
List<ToInstrumentInfo> toInstruments =
filterAndSortDefinitions(definitions, methodInfo.getClassFileLines());
for (ToInstrumentInfo toInstrumentInfo : toInstruments) {
ProbeDefinition definition = toInstrumentInfo.definition;
List<DiagnosticMessage> probeDiagnostics = diagnostics.get(definition.getProbeId());
Expand All @@ -541,7 +543,8 @@ static class ToInstrumentInfo {
}
}

private List<ToInstrumentInfo> filterAndSortDefinitions(List<ProbeDefinition> definitions) {
private List<ToInstrumentInfo> filterAndSortDefinitions(
List<ProbeDefinition> definitions, ClassFileLines classFileLines) {
List<ToInstrumentInfo> toInstrument = new ArrayList<>();
List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
boolean addedExceptionProbe = false;
Expand All @@ -566,7 +569,8 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(List<ProbeDefinition> de
if (!capturedContextProbes.isEmpty()) {
List<ProbeId> probesIds =
capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
ProbeDefinition referenceDefinition = selectReferenceDefinition(capturedContextProbes);
ProbeDefinition referenceDefinition =
selectReferenceDefinition(capturedContextProbes, classFileLines);
toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
}
// LOGGER.debug("exception probe is already instrumented for {}", preciseWhere);
Expand All @@ -581,17 +585,61 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(List<ProbeDefinition> de
}

// Log & Span Decoration probes share the same instrumentor so only one definition should be
// selected to
// generate the instrumentation. Log probes needs capture limits provided by the configuration
// so if the list of definition contains at least 1 log probe this is the log probe that need to
// be picked.
// TODO: handle the conflicting limits for log probes + mixing CaptureSnapshot or not
private ProbeDefinition selectReferenceDefinition(List<ProbeDefinition> capturedContextProbes) {
ProbeDefinition first = capturedContextProbes.get(0);
return capturedContextProbes.stream()
.filter(it -> it instanceof LogProbe)
.findFirst()
.orElse(first);
// used to generate the instrumentation. This method generate a synthetic definition that
// match the type of the probe to instrument: if at least one probe is LogProbe then we are
// creating a LogProbe definition. The synthetic definition contains the union of all the capture,
// snapshot and evaluateAt parameters.
private ProbeDefinition selectReferenceDefinition(
List<ProbeDefinition> capturedContextProbes, ClassFileLines classFileLines) {
boolean hasLogProbe = false;
MethodLocation evaluateAt = MethodLocation.EXIT;
LogProbe.Capture capture = null;
boolean captureSnapshot = false;
Where where = capturedContextProbes.get(0).getWhere();
ProbeId probeId = capturedContextProbes.get(0).getProbeId();
for (ProbeDefinition definition : capturedContextProbes) {
if (definition instanceof LogProbe) {
if (definition instanceof ExceptionProbe) {
where = Where.convertLineToMethod(where, classFileLines);
}
hasLogProbe = true;
LogProbe logProbe = (LogProbe) definition;
captureSnapshot = captureSnapshot | logProbe.isCaptureSnapshot();
capture = mergeCapture(capture, logProbe.getCapture());
}
if (definition.getEvaluateAt() == MethodLocation.ENTRY
|| definition.getEvaluateAt() == MethodLocation.DEFAULT) {
evaluateAt = definition.getEvaluateAt();
}
}
if (hasLogProbe) {
return LogProbe.builder()
.probeId(probeId)
.where(where)
.evaluateAt(evaluateAt)
.capture(capture)
.captureSnapshot(captureSnapshot)
.build();
}
return SpanDecorationProbe.builder()
.probeId(probeId)
.where(where)
.evaluateAt(evaluateAt)
.build();
}

private LogProbe.Capture mergeCapture(LogProbe.Capture current, LogProbe.Capture newCapture) {
if (current == null) {
return newCapture;
}
if (newCapture == null) {
return current;
}
return new LogProbe.Capture(
Math.max(current.getMaxReferenceDepth(), newCapture.getMaxReferenceDepth()),
Math.max(current.getMaxCollectionSize(), newCapture.getMaxCollectionSize()),
Math.max(current.getMaxLength(), newCapture.getMaxLength()),
Math.max(current.getMaxFieldCount(), newCapture.getMaxFieldCount()));
}

private InstrumentationResult.Status preCheckInstrumentation(
Expand Down
5 changes: 5 additions & 0 deletions ...va-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -324,6 +324,11 @@ public void addTag(String tagName, String tagValue) {
public List<Pair<String, String>> getTagsToDecorate() {
return tagsToDecorate;
}

@Override
public boolean isCapturing() {
return true;
}
}

public static SpanDecorationProbe.Builder builder() {
Expand Down
19 changes: 19 additions & 0 deletions ...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -723,6 +723,25 @@ public void fieldExtractorDepth1() throws IOException, URISyntaxException {
assertEquals(DEPTH_REASON, simpleDataFields.get("listValue").getNotCapturedReason());
}

@Test
public void fieldExtractorDuplicateUnionDepth() throws IOException, URISyntaxException {
final String CLASS_NAME = "CapturedSnapshot04";
LogProbe.Builder builder = createProbeBuilder(PROBE_ID, CLASS_NAME, "createSimpleData", "()");
LogProbe probe1 = builder.capture(0, 100, 50, Limits.DEFAULT_FIELD_COUNT).build();
LogProbe probe2 = builder.capture(3, 100, 50, Limits.DEFAULT_FIELD_COUNT).build();
TestSnapshotListener listener = installProbes(CLASS_NAME, probe1, probe2);
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
int result = Reflect.onClass(testClass).call("main", "").get();
assertEquals(143, result);
List<Snapshot> snapshots = assertSnapshots(listener, 2);
CapturedContext.CapturedValue simpleData =
snapshots.get(0).getCaptures().getReturn().getLocals().get("simpleData");
Map<String, CapturedContext.CapturedValue> simpleDataFields = getFields(simpleData);
assertEquals(4, simpleDataFields.size());
assertEquals("foo", simpleDataFields.get("strValue").getValue());
assertEquals(42, simpleDataFields.get("intValue").getValue());
}

@Test
public void fieldExtractorCount2() throws IOException, URISyntaxException {
final String CLASS_NAME = "CapturedSnapshot04";
Expand Down
16 changes: 13 additions & 3 deletions ...gent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -356,25 +356,35 @@ public void ordering() {
.add(metricProbe)
.add(logProbe)
.build();
DebuggerTransformer debuggerTransformer = new DebuggerTransformer(config, configuration);
DebuggerTransformer debuggerTransformer =
new DebuggerTransformer(
config,
configuration,
(definition, result) -> {
if (result.isInstalled()) {
invocationOrder.add(definition);
}
},
new DebuggerSink(
config, new ProbeStatusSink(config, config.getFinalDebuggerSnapshotUrl(), false)));
debuggerTransformer.transform(
ClassLoader.getSystemClassLoader(),
ArrayList.class.getName(), // always FQN
ArrayList.class,
null,
getClassFileBytes(ArrayList.class));
assertEquals(3, invocationOrder.size());
assertEquals(4, invocationOrder.size());
assertEquals(metricProbe, invocationOrder.get(0));
assertEquals(logProbe, invocationOrder.get(1));
assertEquals(spanProbe, invocationOrder.get(2));
assertEquals(spanDecorationProbe, invocationOrder.get(3));
}

<T extends ProbeDefinition> T createMock(
Class<T> clazz, List<ProbeDefinition> invocationOrder, String id) {
ProbeDefinition mock = mock(clazz);
doAnswer(
invocation -> {
invocationOrder.add(mock);
return InstrumentationResult.Status.INSTALLED;
})
.when(mock)
Expand Down
27 changes: 20 additions & 7 deletions ...gger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ public class SpanDecorationProbeInstrumentationTest extends ProbeInstrumentation
private static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f8", 0);
private static final ProbeId PROBE_ID1 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f6", 0);
private static final ProbeId PROBE_ID2 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f7", 0);
private static final ProbeId PROBE_ID3 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f8", 0);
private static final ProbeId PROBE_ID4 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f9", 0);

private TestTraceInterceptor traceInterceptor = new TestTraceInterceptor();

Expand Down Expand Up @@ -315,38 +317,49 @@ public void nullActiveSpan() throws IOException, URISyntaxException {
@Test
public void mixedWithLogProbes() throws IOException, URISyntaxException {
final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
SpanDecorationProbe.Decoration decoration = createDecoration("tag1", "{intLocal}");
SpanDecorationProbe spanDecoProbe =
SpanDecorationProbe.Decoration decoration1 = createDecoration("tag1", "{intLocal}");
SpanDecorationProbe.Decoration decoration2 = createDecoration("tag2", "{arg}");
SpanDecorationProbe spanDecoProbe1 =
createProbeBuilder(
PROBE_ID, ACTIVE, singletonList(decoration), CLASS_NAME, "process", null, null)
PROBE_ID1, ACTIVE, singletonList(decoration1), CLASS_NAME, "process", null, null)
.evaluateAt(MethodLocation.EXIT)
.build();
SpanDecorationProbe spanDecoProbe2 =
createProbeBuilder(
PROBE_ID2, ACTIVE, singletonList(decoration2), CLASS_NAME, "process", null, null)
.evaluateAt(MethodLocation.ENTRY)
.build();
LogProbe logProbe1 =
LogProbe.builder()
.probeId(PROBE_ID1)
.probeId(PROBE_ID3)
.where(CLASS_NAME, "process")
.captureSnapshot(true)
.capture(1, 50, 50, 10)
.build();
LogProbe logProbe2 =
LogProbe.builder()
.probeId(PROBE_ID2)
.probeId(PROBE_ID4)
.where(CLASS_NAME, "process")
.captureSnapshot(true)
.capture(5, 200, 200, 30)
.build();
Configuration configuration =
Configuration.builder()
.setService(SERVICE_NAME)
.add(logProbe1)
.add(logProbe2)
.add(spanDecoProbe)
.add(spanDecoProbe1)
.add(spanDecoProbe2)
.build();
installSpanDecorationProbes(CLASS_NAME, configuration);
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
int result = Reflect.on(testClass).call("main", "1").get();
assertEquals(84, result);
MutableSpan span = traceInterceptor.getFirstSpan();
assertEquals("84", span.getTags().get("tag1"));
assertEquals(PROBE_ID.getId(), span.getTags().get("_dd.di.tag1.probe_id"));
assertEquals(PROBE_ID1.getId(), span.getTags().get("_dd.di.tag1.probe_id"));
assertEquals("1", span.getTags().get("tag2"));
assertEquals(PROBE_ID2.getId(), span.getTags().get("_dd.di.tag2.probe_id"));
List<Snapshot> snapshots = mockSink.getSnapshots();
assertEquals(2, snapshots.size());
CapturedContext.CapturedValue intLocal =
Expand Down

0 comments on commit 3de43de

Please sign in to comment.