Cache entry span code origin information to avoid unnecessary repeated stack walks. (#8029)

evanchooly · web-flow · commit 0d905c6e4067 · 2024-12-02T16:25:54.000+01:00
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/InstrumentationResult.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/InstrumentationResult.java
@@ -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() {
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/MethodInfo.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/MethodInfo.java
@@ -37,4 +37,8 @@ public MethodNode getMethodNode() {
   public ClassFileLines getClassFileLines() {
     return classFileLines;
   }
+
+  public int getMethodStart() {
+    return classFileLines != null ? classFileLines.getMethodStart(methodNode) : -1;
+  }
 }
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java
@@ -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() {
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/ClassFileLines.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/ClassFileLines.java
@@ -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);
   }
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java
@@ -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) {

Original file line number	Diff line number	Diff line change
`@@ -37,4 +37,8 @@ public MethodNode getMethodNode() {`
`37`	`37`	`public ClassFileLines getClassFileLines() {`
`38`	`38`	`return classFileLines;`
`39`	`39`	`}`
	`40`	`+`
	`41`	`+ public int getMethodStart() {`
	`42`	`+ return classFileLines != null ? classFileLines.getMethodStart(methodNode) : -1;`
	`43`	`+ }`
`40`	`44`	`}`