Skip to content

Commit

Permalink
Fix CodeOrigin for @trace annotation
Browse files Browse the repository at this point in the history
add smoke test
  • Loading branch information
jpbempel committed Feb 6, 2025
1 parent 9ec7fde commit 683c682
Show file tree
Hide file tree
Showing 7 changed files with 105 additions and 26 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ public interface ExceptionDebugger {
public interface CodeOriginRecorder {
String captureCodeOrigin(boolean entry);

String captureCodeOrigin(Method method, boolean entry);
String captureCodeOrigin(Method method, boolean entry, boolean instrument);
}

private static volatile ProbeResolver probeResolver;
Expand Down Expand Up @@ -405,10 +405,14 @@ public static String captureCodeOrigin(boolean entry) {
}

public static String captureCodeOrigin(Method method, boolean entry) {
return captureCodeOrigin(method, entry, true);
}

public static String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
try {
CodeOriginRecorder recorder = codeOriginRecorder;
if (recorder != null) {
return recorder.captureCodeOrigin(method, entry);
return recorder.captureCodeOrigin(method, entry, instrument);
}
} catch (Exception ex) {
LOGGER.debug("Error in captureCodeOrigin: ", ex);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,20 +71,25 @@ public String captureCodeOrigin(boolean entry) {
element.getMethodName(),
null,
String.valueOf(element.getLineNumber()));
probe = createProbe(fingerprint, entry, where);
probe = createProbe(fingerprint, entry, where, true);

LOG.debug("Creating probe for location {}", where);
}
return probe.getId();
}

@Override
public String captureCodeOrigin(Method method, boolean entry) {
public String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
String fingerprint = method.toString();
CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
if (probe == null) {
probe = createProbe(fingerprint, entry, Where.of(method));
probe = createProbe(fingerprint, entry, Where.of(method), instrument);
LOG.debug("Creating probe for method {}", fingerprint);
} else if (!instrument) {
// direct call to fill code origin info without using probe instrumentation
// buildLocation should be called before in order to gather location info
probe.commit(
CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
}
return probe.getId();
}
Expand All @@ -105,11 +110,12 @@ public void registerLogProbe(CodeOriginProbe probe) {
.build());
}

private CodeOriginProbe createProbe(String fingerPrint, boolean entry, Where where) {
private CodeOriginProbe createProbe(
String fingerPrint, boolean entry, Where where, boolean instrument) {
CodeOriginProbe probe;
AgentSpan span = AgentTracer.activeSpan();

probe = new CodeOriginProbe(ProbeId.newId(), entry, where);
probe = new CodeOriginProbe(ProbeId.newId(), entry, where, instrument);
addFingerprint(fingerPrint, probe);
CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,19 +27,23 @@
public class CodeOriginProbe extends ProbeDefinition {
private static final Logger LOGGER = LoggerFactory.getLogger(CodeOriginProbe.class);

private final boolean instrument;
private final boolean entrySpanProbe;

private String signature;

public CodeOriginProbe(ProbeId probeId, boolean entry, Where where) {
public CodeOriginProbe(ProbeId probeId, boolean entry, Where where, boolean instrument) {
super(LANGUAGE, probeId, (Tag[]) null, where, MethodLocation.ENTRY);
this.instrument = instrument;
this.entrySpanProbe = entry;
}

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

@Override
Expand All @@ -55,6 +59,10 @@ public void commit(
List<AgentSpan> agentSpans =
entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);

if (location == null) {
LOGGER.debug("Code origin probe {} has no location", id);
return;
}
for (AgentSpan s : agentSpans) {
if (s.getTag(DD_CODE_ORIGIN_TYPE) == null) {
s.setTag(DD_CODE_ORIGIN_TYPE, entrySpanProbe ? "entry" : "exit");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -105,8 +105,8 @@ public void withDebug1() throws Exception {
final String className = "com.datadog.debugger.CodeOrigin02";
installProbes();
final Class<?> testClass = compileAndLoadClass(className);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true, true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false, true);
checkResults(testClass, "fullTrace", 0);
checkResults(testClass, "debug_1", 2);
}
Expand All @@ -117,8 +117,8 @@ public void withLogProbe() throws Exception {
installProbes(
createProbeBuilder(PROBE_ID, CLASS_NAME, "entry", "()").captureSnapshot(true).build());
final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true, true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false, true);
checkResults(testClass, "debug_1", 3);
}

Expand All @@ -127,10 +127,13 @@ public void doubleEntry() throws IOException, URISyntaxException {
final String className = "com.datadog.debugger.CodeOrigin05";

installProbes(
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(className, "entry", "()", "53")),
new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62")),
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(className, "entry", "()", "53"), true),
new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62"), true),
new CodeOriginProbe(
CODE_ORIGIN_DOUBLE_ENTRY_ID, true, Where.of(className, "doubleEntry", "()", "66")));
CODE_ORIGIN_DOUBLE_ENTRY_ID,
true,
Where.of(className, "doubleEntry", "()", "66"),
true));
final Class<?> testClass = compileAndLoadClass(className);
checkResults(testClass, "fullTrace", 0);
List<? extends MutableSpan> trace = traceInterceptor.getTrace();
Expand All @@ -143,7 +146,7 @@ public void doubleEntry() throws IOException, URISyntaxException {
public void stackDepth() throws IOException, URISyntaxException {
final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
installProbes(
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(CLASS_NAME, "exit", "()", "39")));
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(CLASS_NAME, "exit", "()", "39"), true));

Class<?> testClass = compileAndLoadClass("com.datadog.debugger.CodeOrigin04");
countFrames(testClass, 10);
Expand Down Expand Up @@ -187,7 +190,8 @@ public void testCaptureCodeOriginWithExplicitInfo()
installProbes();
CodeOriginProbe probe =
codeOriginRecorder.getProbe(
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true));
codeOriginRecorder.captureCodeOrigin(
testClass.getMethod("main", int.class), true, true));
assertNotNull(probe, "The probe should have been created.");
assertTrue(probe.entrySpanProbe(), "Should be an entry probe.");
}
Expand All @@ -199,18 +203,18 @@ public void testDuplicateInstrumentations()
final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
installProbes();
String probe1 =
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true, true);
String probe2 =
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true, true);
assertEquals(probe1, probe2);
}

@NotNull
private CodeOriginProbe[] codeOriginProbes(String type) {
CodeOriginProbe entry =
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(type, "entry", "()", "53"));
new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(type, "entry", "()", "53"), true);
CodeOriginProbe exit =
new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(type, "exit", "()", "60"));
new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(type, "exit", "()", "60"), true);
return new CodeOriginProbe[] {entry, exit};
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECORATE;

import datadog.trace.bootstrap.debugger.DebuggerContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import java.lang.invoke.MethodType;
import java.lang.reflect.Method;
Expand All @@ -13,7 +14,9 @@ public class TraceAdvice {

@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope onEnter(@Advice.Origin final Method method) {
return activateSpan(DECORATE.startMethodSpan(method));
AgentScope agentScope = activateSpan(DECORATE.startMethodSpan(method));
DebuggerContext.captureCodeOrigin(method, true, false);
return agentScope;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import datadog.trace.api.InstrumenterConfig;
import datadog.trace.api.Trace;
import datadog.trace.bootstrap.debugger.spanorigin.CodeOriginInfo;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
import datadog.trace.bootstrap.instrumentation.decorator.AsyncResultDecorator;
Expand Down Expand Up @@ -91,7 +90,6 @@ public AgentSpan startMethodSpan(Method method) {
afterStart(span);
span.setResourceName(resourceName);

CodeOriginInfo.entry(method);
if (measured || InstrumenterConfig.get().isMethodMeasured(method)) {
span.setMeasured(true);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
package datadog.smoketest;

import com.datadog.debugger.probe.SpanDecorationProbe;
import datadog.trace.api.DDTags;
import datadog.trace.test.agent.decoder.DecodedSpan;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;

import java.nio.file.Path;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

import static org.junit.jupiter.api.Assertions.assertEquals;

public class CodeOriginIntegrationTest extends ServerAppDebuggerIntegrationTest {

private static final String DD_CODE_ORIGIN_FRAMES_0_FILE = String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "file");
private static final String DD_CODE_ORIGIN_FRAMES_0_METHOD = String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "method");
private static final String DD_CODE_ORIGIN_FRAMES_0_SIGNATURE = String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "signature");
private static final String DD_CODE_ORIGIN_FRAMES_0_LINE = String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "line");

@Override
protected ProcessBuilder createProcessBuilder(Path logFilePath, String... params) {
List<String> commandParams = getDebuggerCommandParams();
commandParams.add("-Ddd.trace.enabled=true"); // explicitly enable tracer
commandParams.add("-Ddd.code.origin.for.spans.enabled=true");
return ProcessBuilderHelper.createProcessBuilder(
commandParams, logFilePath, getAppClass(), params);
}

@Test
@DisplayName("testCodeOriginTraceAnnotation")
void testCodeOriginTraceAnnotation() throws Exception {
execute(appUrl, TRACED_METHOD_NAME);
waitForInstrumentation(appUrl);
execute(appUrl, TRACED_METHOD_NAME);
AtomicBoolean codeOrigin = new AtomicBoolean();
registerTraceListener(
decodedTrace -> {
for (DecodedSpan span : decodedTrace.getSpans()) {
if (isTracedFullMethodSpan(span)) {
if (span.getMeta().containsKey(DDTags.DD_CODE_ORIGIN_TYPE)) {
assertEquals("entry", span.getMeta().get(DDTags.DD_CODE_ORIGIN_TYPE));
assertEquals("ServerDebuggerTestApplication.java", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_FILE));
assertEquals("runTracedMethod", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_METHOD));
assertEquals("(java.lang.String)", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_SIGNATURE));
assertEquals("133", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_LINE));
codeOrigin.set(true);
}
}
}
});
processRequests(codeOrigin::get);
}
}

0 comments on commit 683c682

Please sign in to comment.