Skip to content

Commit

Permalink
Merge pull request #6444 from DataDog/jpbempel/probe-order
Browse files Browse the repository at this point in the history 
Enforce ordering of probe instrumentation
  • Loading branch information
@jpbempel
jpbempel authored Jan 12, 2024
2 parents f471f93 + f46ecf4 commit 027316b
Show file tree
Hide file tree
Showing 11 changed files with 230 additions and 64 deletions.
5 changes: 4 additions & 1 deletion ...er/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,11 @@ public interface ValueSerializer {
private static volatile Tracer tracer;
private static volatile ValueSerializer valueSerializer;

public static void init(ProbeResolver probeResolver, MetricForwarder metricForwarder) {
public static void initProbeResolver(ProbeResolver probeResolver) {
DebuggerContext.probeResolver = probeResolver;
}

public static void initMetricForwarder(MetricForwarder metricForwarder) {
DebuggerContext.metricForwarder = metricForwarder;
}

Expand Down
6 changes: 4 additions & 2 deletions dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,8 @@ public static synchronized void run(
sink = debuggerSink;
StatsdMetricForwarder statsdMetricForwarder =
new StatsdMetricForwarder(config, probeStatusSink);
DebuggerContext.init(configurationUpdater, statsdMetricForwarder);
DebuggerContext.initProbeResolver(configurationUpdater);
DebuggerContext.initMetricForwarder(statsdMetricForwarder);
DebuggerContext.initClassFilter(new DenyListHelper(null)); // default hard coded deny list
snapshotSerializer = new JsonSnapshotSerializer();
DebuggerContext.initValueSerializer(snapshotSerializer);
Expand Down Expand Up @@ -168,7 +169,8 @@ static ClassFileTransformer setupInstrumentTheWorldTransformer(
LOGGER.info("install Instrument-The-World transformer");
DebuggerTransformer transformer =
createTransformer(config, Configuration.builder().build(), null, debuggerSink);
DebuggerContext.init(transformer::instrumentTheWorldResolver, statsdMetricForwarder);
DebuggerContext.initProbeResolver(transformer::instrumentTheWorldResolver);
DebuggerContext.initMetricForwarder(statsdMetricForwarder);
instrumentation.addTransformer(transformer);
return transformer;
}
Expand Down
68 changes: 49 additions & 19 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 @@ -6,8 +6,10 @@
import com.datadog.debugger.instrumentation.DiagnosticMessage;
import com.datadog.debugger.instrumentation.InstrumentationResult;
import com.datadog.debugger.probe.LogProbe;
import com.datadog.debugger.probe.MetricProbe;
import com.datadog.debugger.probe.ProbeDefinition;
import com.datadog.debugger.probe.SpanDecorationProbe;
import com.datadog.debugger.probe.SpanProbe;
import com.datadog.debugger.probe.Where;
import com.datadog.debugger.sink.DebuggerSink;
import com.datadog.debugger.sink.ProbeStatusSink;
Expand All @@ -30,6 +32,7 @@
import java.security.CodeSource;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
Expand Down Expand Up @@ -66,6 +69,8 @@ public class DebuggerTransformer implements ClassFileTransformer {
private static final String INSTRUMENTATION_FAILS = "Instrumentation fails for %s";
private static final String CANNOT_FIND_LINE = "No executable code was found at %s:L%s";
private static final Pattern COMMA_PATTERN = Pattern.compile(",");
private static final List<Class<?>> PROBE_ORDER =
Arrays.asList(MetricProbe.class, LogProbe.class, SpanDecorationProbe.class, SpanProbe.class);

private final Config config;
private final TransformerDefinitionMatcher definitionMatcher;
Expand Down Expand Up @@ -505,26 +510,13 @@ private InstrumentationResult applyInstrumentation(
preCheckInstrumentation(diagnostics, classLoader, methodNode);
if (status != InstrumentationResult.Status.ERROR) {
try {
List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
for (ProbeDefinition definition : definitions) {
// Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
// and therefore need to be instrumented once
if (definition instanceof LogProbe || definition instanceof SpanDecorationProbe) {
capturedContextProbes.add(definition);
} else {
List<DiagnosticMessage> probeDiagnostics = diagnostics.get(definition.getProbeId());
status = definition.instrument(classLoader, classNode, methodNode, probeDiagnostics);
}
}
if (capturedContextProbes.size() > 0) {
List<ProbeId> probesIds =
capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
ProbeDefinition referenceDefinition = selectReferenceDefinition(capturedContextProbes);
List<DiagnosticMessage> probeDiagnostics =
diagnostics.get(referenceDefinition.getProbeId());
List<ToInstrumentInfo> toInstruments = filterAndSortDefinitions(definitions);
for (ToInstrumentInfo toInstrumentInfo : toInstruments) {
ProbeDefinition definition = toInstrumentInfo.definition;
List<DiagnosticMessage> probeDiagnostics = diagnostics.get(definition.getProbeId());
status =
referenceDefinition.instrument(
classLoader, classNode, methodNode, probeDiagnostics, probesIds);
definition.instrument(
classLoader, classNode, methodNode, probeDiagnostics, toInstrumentInfo.probeIds);
}
} catch (Throwable t) {
log.warn("Exception during instrumentation: ", t);
Expand All @@ -536,6 +528,44 @@ private InstrumentationResult applyInstrumentation(
return new InstrumentationResult(status, diagnostics, classNode, methodNode);
}

static class ToInstrumentInfo {
final ProbeDefinition definition;
final List<ProbeId> probeIds;

ToInstrumentInfo(ProbeDefinition definition, List<ProbeId> probeIds) {
this.definition = definition;
this.probeIds = probeIds;
}
}

private List<ToInstrumentInfo> filterAndSortDefinitions(List<ProbeDefinition> definitions) {
List<ToInstrumentInfo> toInstrument = new ArrayList<>();
List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
for (ProbeDefinition definition : definitions) {
// Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
// and therefore need to be instrumented once
if (definition instanceof LogProbe || definition instanceof SpanDecorationProbe) {
capturedContextProbes.add(definition);
} else {
toInstrument.add(new ToInstrumentInfo(definition, singletonList(definition.getProbeId())));
}
}
if (!capturedContextProbes.isEmpty()) {
List<ProbeId> probesIds =
capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
ProbeDefinition referenceDefinition = selectReferenceDefinition(capturedContextProbes);
toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
}
// ordering: metric < log < span decoration < span
toInstrument.sort(
(info1, info2) -> {
int idx1 = PROBE_ORDER.indexOf(info1.definition.getClass());
int idx2 = PROBE_ORDER.indexOf(info2.definition.getClass());
return Integer.compare(idx1, idx2);
});
return toInstrument;
}

// 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
Expand Down
10 changes: 0 additions & 10 deletions dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ProbeDefinition.java
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
package com.datadog.debugger.probe;

import static java.util.Collections.singletonList;

import com.datadog.debugger.agent.Generated;
import com.datadog.debugger.instrumentation.DiagnosticMessage;
import com.datadog.debugger.instrumentation.InstrumentationResult;
Expand Down Expand Up @@ -127,14 +125,6 @@ private static void initTagMap(Map<String, String> tagMap, Tag[] tags) {
}
}

public InstrumentationResult.Status instrument(
ClassLoader classLoader,
ClassNode classNode,
MethodNode methodNode,
List<DiagnosticMessage> diagnostics) {
return instrument(classLoader, classNode, methodNode, diagnostics, singletonList(getProbeId()));
}

public abstract InstrumentationResult.Status instrument(
ClassLoader classLoader,
ClassNode classNode,
Expand Down
112 changes: 101 additions & 11 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 @@ -8,6 +8,7 @@
import static com.datadog.debugger.util.MoshiSnapshotHelper.REDACTED_TYPE_REASON;
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
import static datadog.trace.bootstrap.debugger.util.Redaction.REDACTED_VALUE;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
Expand All @@ -29,19 +30,27 @@
import com.datadog.debugger.el.values.StringValue;
import com.datadog.debugger.instrumentation.InstrumentationResult;
import com.datadog.debugger.probe.LogProbe;
import com.datadog.debugger.probe.MetricProbe;
import com.datadog.debugger.probe.ProbeDefinition;
import com.datadog.debugger.probe.SpanDecorationProbe;
import com.datadog.debugger.probe.SpanProbe;
import com.datadog.debugger.probe.Where;
import com.datadog.debugger.sink.DebuggerSink;
import com.datadog.debugger.sink.ProbeStatusSink;
import com.datadog.debugger.sink.Snapshot;
import com.datadog.debugger.util.MoshiHelper;
import com.datadog.debugger.util.MoshiSnapshotTestHelper;
import com.datadog.debugger.util.SerializerWithLimits;
import com.squareup.moshi.JsonAdapter;
import datadog.trace.agent.tooling.TracerInstaller;
import datadog.trace.api.Config;
import datadog.trace.api.interceptor.MutableSpan;
import datadog.trace.api.sampling.Sampler;
import datadog.trace.bootstrap.debugger.*;
import datadog.trace.bootstrap.debugger.el.ValueReferences;
import datadog.trace.bootstrap.debugger.util.Redaction;
import datadog.trace.core.CoreTracer;
import datadog.trace.core.DDSpan;
import groovy.lang.GroovyClassLoader;
import java.io.File;
import java.io.IOException;
Expand Down Expand Up @@ -87,7 +96,7 @@

public class CapturedSnapshotTest {
private static final String LANGUAGE = "java";
private static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f8", 0);
private static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f5", 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);
Expand Down Expand Up @@ -171,7 +180,7 @@ public void resolutionFails() throws IOException, URISyntaxException {
DebuggerTransformerTest.TestSnapshotListener listener =
installSingleProbe(CLASS_NAME, "main", "int (java.lang.String)", "8");
DebuggerAgentHelper.injectSink(listener);
DebuggerContext.init((encodedProbeId) -> null, null);
DebuggerContext.initProbeResolver((encodedProbeId) -> null);
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
int result = Reflect.on(testClass).call("main", "1").get();
assertEquals(3, result);
Expand All @@ -186,11 +195,10 @@ public void resolutionThrows() throws IOException, URISyntaxException {
DebuggerTransformerTest.TestSnapshotListener listener =
installProbes(CLASS_NAME, lineProbe, methodProbe);
DebuggerAgentHelper.injectSink(listener);
DebuggerContext.init(
DebuggerContext.initProbeResolver(
(encodedProbeId) -> {
throw new IllegalArgumentException("oops");
},
null);
});
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
int result = Reflect.on(testClass).call("main", "1").get();
assertEquals(3, result);
Expand Down Expand Up @@ -2038,6 +2046,75 @@ public void lineRecord() throws IOException, URISyntaxException {
assertCaptureFields(capturedContext, "age", Integer.TYPE.getTypeName(), "0");
}

@Test
public void allProbesSameMethod() throws IOException, URISyntaxException {
final String CLASS_NAME = "CapturedSnapshot01";
final String METRIC_NAME = "count";
Where where = new Where().typeName(CLASS_NAME).methodName("main");
Configuration configuration =
Configuration.builder()
.add(
SpanDecorationProbe.builder()
.probeId(PROBE_ID)
.where(where)
.targetSpan(SpanDecorationProbe.TargetSpan.ACTIVE)
.decorate(
new SpanDecorationProbe.Decoration(
null,
Arrays.asList(
new SpanDecorationProbe.Tag(
"tag1",
new SpanDecorationProbe.TagValue(
"value1", parseTemplate("value1"))))))
.build())
.add(SpanProbe.builder().probeId(PROBE_ID1).where(where).build())
.add(
MetricProbe.builder()
.probeId(PROBE_ID2)
.metricName(METRIC_NAME)
.kind(MetricProbe.MetricKind.COUNT)
.where(where)
.build())
.add(LogProbe.builder().probeId(PROBE_ID3).where(where).build())
.build();

CoreTracer tracer = CoreTracer.builder().build();
TracerInstaller.forceInstallGlobalTracer(tracer);
SpanDecorationProbeInstrumentationTest.TestTraceInterceptor traceInterceptor =
new SpanDecorationProbeInstrumentationTest.TestTraceInterceptor();
tracer.addTraceInterceptor(traceInterceptor);
try {
DebuggerTransformerTest.TestSnapshotListener snapshotListener =
installProbes(CLASS_NAME, configuration);
DebuggerContext.initTracer(new DebuggerTracer(mock(ProbeStatusSink.class)));
MetricProbesInstrumentationTest.MetricForwarderListener metricListener =
new MetricProbesInstrumentationTest.MetricForwarderListener();
DebuggerContext.initMetricForwarder(metricListener);
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
int result = Reflect.on(testClass).call("main", "1").get();
// log probe
assertEquals(3, result);
assertEquals(1, snapshotListener.snapshots.size());
Snapshot snapshot = snapshotListener.snapshots.get(0);
assertEquals(PROBE_ID3.getId(), snapshot.getProbe().getId());
// span (deco) probe
assertEquals(1, traceInterceptor.getTrace().size());
MutableSpan span = traceInterceptor.getFirstSpan();
assertEquals(CLASS_NAME + ".main", span.getResourceName());
assertEquals(PROBE_ID1.getId(), span.getTag("debugger.probeid"));
assertEquals("value1", span.getTag("tag1"));
// correlation between log and created span
assertEquals(snapshot.getSpanId(), String.valueOf(((DDSpan) span).getSpanId()));
// metric probe
assertTrue(metricListener.counters.containsKey(METRIC_NAME));
assertEquals(1, metricListener.counters.get(METRIC_NAME).longValue());
assertArrayEquals(
new String[] {"debugger.probeid:" + PROBE_ID2.getId()}, metricListener.lastTags);
} finally {
TracerInstaller.forceInstallGlobalTracer(null);
}
}

private DebuggerTransformerTest.TestSnapshotListener setupInstrumentTheWorldTransformer(
String excludeFileName) {
Config config = mock(Config.class);
Expand Down Expand Up @@ -2114,13 +2191,18 @@ private DebuggerTransformerTest.TestSnapshotListener installProbes(
new DebuggerTransformer(config, configuration, instrumentationListener, listener);
instr.addTransformer(currentTransformer);
DebuggerAgentHelper.injectSink(listener);
DebuggerContext.init((encodedId) -> resolver(encodedId, logProbes), null);
DebuggerContext.initProbeResolver(
(encodedId) -> resolver(encodedId, logProbes, configuration.getSpanDecorationProbes()));
DebuggerContext.initClassFilter(new DenyListHelper(null));
DebuggerContext.initValueSerializer(new JsonSnapshotSerializer());
for (LogProbe probe : logProbes) {
if (probe.getSampling() != null) {
ProbeRateLimiter.setRate(
probe.getId(), probe.getSampling().getSnapshotsPerSecond(), probe.isCaptureSnapshot());
if (logProbes != null) {
for (LogProbe probe : logProbes) {
if (probe.getSampling() != null) {
ProbeRateLimiter.setRate(
probe.getId(),
probe.getSampling().getSnapshotsPerSecond(),
probe.isCaptureSnapshot());
}
}
}
if (configuration.getSampling() != null) {
Expand All @@ -2129,12 +2211,20 @@ private DebuggerTransformerTest.TestSnapshotListener installProbes(
return listener;
}

private ProbeImplementation resolver(String encodedId, Collection<LogProbe> logProbes) {
private ProbeImplementation resolver(
String encodedId,
Collection<LogProbe> logProbes,
Collection<SpanDecorationProbe> spanDecorationProbes) {
for (LogProbe probe : logProbes) {
if (probe.getProbeId().getEncodedId().equals(encodedId)) {
return probe;
}
}
for (SpanDecorationProbe probe : spanDecorationProbes) {
if (probe.getProbeId().getEncodedId().equals(encodedId)) {
return probe;
}
}
return null;
}

Expand Down
1 change: 1 addition & 0 deletions ...ava-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTracerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ public void setError() {

@Test
public void noApi() {
AgentTracer.forceRegister(null);
ProbeStatusSink probeStatusSink = mock(ProbeStatusSink.class);
DebuggerTracer debuggerTracer = new DebuggerTracer(probeStatusSink);
DebuggerSpan span =
Expand Down
Loading

0 comments on commit 027316b

Please sign in to comment.