From 3de43de7442c28644cc16d6c4d665f0730348b61 Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bempel Date: Tue, 9 Jul 2024 15:55:54 +0200 Subject: [PATCH] Fix mixing log/span decoration probes (#7246) 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. --- .../debugger/agent/DebuggerTransformer.java | 76 +++++++++++++++---- .../debugger/probe/SpanDecorationProbe.java | 5 ++ .../debugger/agent/CapturedSnapshotTest.java | 19 +++++ .../agent/DebuggerTransformerTest.java | 16 +++- ...panDecorationProbeInstrumentationTest.java | 27 +++++-- 5 files changed, 119 insertions(+), 24 deletions(-) diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java index 4b75f6f8586..290bddbd0f7 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java +++ b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java @@ -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; @@ -515,7 +516,8 @@ private InstrumentationResult applyInstrumentation( InstrumentationResult.Status status = preCheckInstrumentation(diagnostics, methodInfo); if (status != InstrumentationResult.Status.ERROR) { try { - List toInstruments = filterAndSortDefinitions(definitions); + List toInstruments = + filterAndSortDefinitions(definitions, methodInfo.getClassFileLines()); for (ToInstrumentInfo toInstrumentInfo : toInstruments) { ProbeDefinition definition = toInstrumentInfo.definition; List probeDiagnostics = diagnostics.get(definition.getProbeId()); @@ -541,7 +543,8 @@ static class ToInstrumentInfo { } } - private List filterAndSortDefinitions(List definitions) { + private List filterAndSortDefinitions( + List definitions, ClassFileLines classFileLines) { List toInstrument = new ArrayList<>(); List capturedContextProbes = new ArrayList<>(); boolean addedExceptionProbe = false; @@ -566,7 +569,8 @@ private List filterAndSortDefinitions(List de if (!capturedContextProbes.isEmpty()) { List 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); @@ -581,17 +585,61 @@ private List filterAndSortDefinitions(List 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 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 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( diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java index 3ee3a1b4746..546d09194b6 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java +++ b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java @@ -324,6 +324,11 @@ public void addTag(String tagName, String tagValue) { public List> getTagsToDecorate() { return tagsToDecorate; } + + @Override + public boolean isCapturing() { + return true; + } } public static SpanDecorationProbe.Builder builder() { diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java index 099b6a28012..5534ca2212d 100644 --- a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java +++ b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java @@ -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 snapshots = assertSnapshots(listener, 2); + CapturedContext.CapturedValue simpleData = + snapshots.get(0).getCaptures().getReturn().getLocals().get("simpleData"); + Map 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"; diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java index 37ad4575eb4..78b4aeecd91 100644 --- a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java +++ b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java @@ -356,17 +356,28 @@ 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 createMock( @@ -374,7 +385,6 @@ T createMock( ProbeDefinition mock = mock(clazz); doAnswer( invocation -> { - invocationOrder.add(mock); return InstrumentationResult.Status.INSTALLED; }) .when(mock) diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java index a3d06d6064c..3c249fcbd36 100644 --- a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java +++ b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java @@ -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(); @@ -315,30 +317,39 @@ 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); @@ -346,7 +357,9 @@ PROBE_ID, ACTIVE, singletonList(decoration), CLASS_NAME, "process", null, null) 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 snapshots = mockSink.getSnapshots(); assertEquals(2, snapshots.size()); CapturedContext.CapturedValue intLocal =