Skip to content

Commit

Permalink
Cache entry span code origin information to avoid unnecessary repeate…
Browse files Browse the repository at this point in the history
…d stack walks. (#8029)
  • Loading branch information
evanchooly authored Dec 2, 2024
1 parent d7a0014 commit 0d905c6
Showing 5 changed files with 89 additions and 11 deletions.
Original file line number Diff line number Diff line change
@@ -18,8 +18,10 @@ public enum Status {

private final Status status;
private final Map<ProbeId, List<DiagnosticMessage>> diagnostics;
private String typeName;
private String methodName;
private final String sourceFileName;
private final String typeName;
private final String methodName;
private final int methodStart;

public static class Factory {
public static InstrumentationResult blocked(String className) {
@@ -41,8 +43,10 @@ public InstrumentationResult(
this(
status,
diagnostics,
methodInfo.getClassNode().sourceFile,
methodInfo.getClassNode().name.replace('/', '.'),
methodInfo.getMethodNode().name);
methodInfo.getMethodNode().name,
methodInfo.getMethodStart());
}

public InstrumentationResult(
@@ -54,6 +58,23 @@ public InstrumentationResult(
this.diagnostics = diagnostics;
this.typeName = className;
this.methodName = methodName;
this.methodStart = -1;
this.sourceFileName = null;
}

public InstrumentationResult(
Status status,
Map<ProbeId, List<DiagnosticMessage>> diagnostics,
String sourceFileName,
String className,
String methodName,
int methodStart) {
this.status = status;
this.diagnostics = diagnostics;
this.sourceFileName = sourceFileName;
this.typeName = className;
this.methodName = methodName;
this.methodStart = methodStart;
}

public boolean isError() {
@@ -80,6 +101,14 @@ public String getMethodName() {
return methodName;
}

public int getMethodStart() {
return methodStart;
}

public String getSourceFileName() {
return sourceFileName;
}

@Generated
@Override
public String toString() {
Original file line number Diff line number Diff line change
@@ -37,4 +37,8 @@ public MethodNode getMethodNode() {
public ClassFileLines getClassFileLines() {
return classFileLines;
}

public int getMethodStart() {
return classFileLines != null ? classFileLines.getMethodStart(methodNode) : -1;
}
}
Original file line number Diff line number Diff line change
@@ -5,6 +5,7 @@
import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_TYPE;
import static java.lang.String.format;
import static java.util.Arrays.asList;
import static java.util.Collections.emptyList;
import static java.util.Collections.singletonList;

import com.datadog.debugger.agent.DebuggerAgent;
@@ -31,6 +32,8 @@ public class CodeOriginProbe extends LogProbe implements ForceMethodInstrumentat

private final boolean entrySpanProbe;

private List<StackTraceElement> stackTraceElements;

public CodeOriginProbe(ProbeId probeId, boolean entry, Where where, int maxFrames) {
super(LANGUAGE, probeId, null, where, MethodLocation.EXIT, null, null, true, null, null, null);
this.entrySpanProbe = entry;
@@ -82,8 +85,25 @@ public void commit(
span.getLocalRootSpan().setTag(getId(), (String) null); // clear possible span reference
}

private List<StackTraceElement> findLocation() {
if (entrySpanProbe && stackTraceElements == null) {
ProbeLocation probeLocation = getLocation();
List<String> lines = probeLocation.getLines();
int line = lines == null ? -1 : Integer.parseInt(lines.get(0));
stackTraceElements =
singletonList(
new StackTraceElement(
probeLocation.getType(),
probeLocation.getMethod(),
probeLocation.getFile(),
line));
}
return stackTraceElements != null ? stackTraceElements : emptyList();
}

private void applySpanOriginTags(AgentSpan span, String snapshotId) {
List<StackTraceElement> entries = getUserStackFrames();
List<StackTraceElement> entries =
stackTraceElements != null ? stackTraceElements : getUserStackFrames();
List<AgentSpan> agentSpans =
entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);

@@ -92,7 +112,8 @@ private void applySpanOriginTags(AgentSpan span, String snapshotId) {

for (int i = 0; i < entries.size(); i++) {
StackTraceElement info = entries.get(i);
s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "file"), info.getFileName());
String fileName = info.getFileName();
s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "file"), fileName);
s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "method"), info.getMethodName());
s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "line"), info.getLineNumber());
s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "type"), info.getClassName());
@@ -127,12 +148,25 @@ private AgentSpan findSpan(AgentSpan candidate) {
public void buildLocation(InstrumentationResult result) {
String type = where.getTypeName();
String method = where.getMethodName();
List<String> lines = null;

String file = where.getSourceFile();

if (result != null) {
type = result.getTypeName();
method = result.getMethodName();
if (result.getMethodStart() != -1) {
lines = singletonList(String.valueOf(result.getMethodStart()));
}
if (file == null) {
file = result.getSourceFileName();
}
if (entrySpanProbe) {
stackTraceElements =
singletonList(new StackTraceElement(type, method, file, result.getMethodStart()));
}
}
// drop line number for code origin probe
this.location = new ProbeLocation(type, method, where.getSourceFile(), null);
this.location = new ProbeLocation(type, method, file, lines);
}

private List<StackTraceElement> getUserStackFrames() {
Original file line number Diff line number Diff line change
@@ -13,6 +13,7 @@

public class ClassFileLines {
private final Map<Integer, List<MethodNode>> methodByLine = new HashMap<>();
private final Map<String, Integer> methodStarts = new HashMap<>();
private final TreeMap<Integer, LabelNode> lineLabels = new TreeMap<>();

public ClassFileLines(ClassNode classNode) {
@@ -21,6 +22,7 @@ public ClassFileLines(ClassNode classNode) {
while (currentNode != null) {
if (currentNode.getType() == AbstractInsnNode.LINE) {
LineNumberNode lineNode = (LineNumberNode) currentNode;
methodStarts.putIfAbsent(methodNode.name + methodNode.desc, lineNode.line);
// on the same line, we can have multiple methods (lambdas, inner classes, etc)
List<MethodNode> methodNodes =
methodByLine.computeIfAbsent(lineNode.line, k -> new ArrayList<>());
@@ -40,6 +42,10 @@ public ClassFileLines(ClassNode classNode) {
}
}

public int getMethodStart(MethodNode node) {
return methodStarts.getOrDefault(node.name + node.desc, -1);
}

public List<MethodNode> getMethodsByLine(int line) {
return methodByLine.get(line);
}
Original file line number Diff line number Diff line change
@@ -8,6 +8,7 @@
import static java.util.Arrays.asList;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -234,6 +235,7 @@ private static void checkEntrySpanTags(MutableSpan span, boolean includeSnapshot
assertKeyPresent(span, DD_CODE_ORIGIN_TYPE);
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "file"));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "line"));
assertNotEquals(-1, span.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "line")));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "method"));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "signature"));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "type"));
@@ -263,6 +265,7 @@ private static void checkExitSpanTags(MutableSpan span, boolean includeSnapshot)
assertKeyPresent(span, DD_CODE_ORIGIN_TYPE);
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "file"));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "line"));
assertNotEquals(-1, span.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "line")));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "method"));
assertKeyPresent(span, format(DD_CODE_ORIGIN_FRAME, 0, "type"));
if (includeSnapshot) {
@@ -271,10 +274,12 @@ private static void checkExitSpanTags(MutableSpan span, boolean includeSnapshot)

MutableSpan rootSpan = span.getLocalRootSpan();
assertEquals(rootSpan.getTag(DD_CODE_ORIGIN_TYPE), "entry", keys);
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 1, "file")));
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 1, "line")));
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 1, "method")));
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 1, "type")));
Object file = rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "file"));
assertNotNull(file, rootSpan.getTags().toString());
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "line")));
assertNotEquals(-1, rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "line")));
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "method")));
assertNotNull(rootSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "type")));
}

private static Set<String> ldKeys(MutableSpan span) {

0 comments on commit 0d905c6

Please sign in to comment.