Fix line probe in method with inline lambdas (#7099)

Multiple methods can match a line if inline lambdas are defined on the same line. we need to collect all methods associated to a line and pick one. Here the first which will be the enclosing method ClassFileLines now collect all methods for a same line
DataDog · May 31, 2024 · dadbd53 · dadbd53
1 parent 7b17119
commit dadbd53
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 28 deletions.
diff --git a/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java b/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
@@ -655,10 +655,15 @@ private MethodNode matchSourceFile(
     }
     Where.SourceLine sourceLine = lines[0]; // assume only 1 range
     int matchingLine = sourceLine.getFrom();
-    MethodNode matchingMethod = classFileLines.getMethodByLine(matchingLine);
-    if (matchingMethod != null) {
-      log.debug("Found lineNode {} method: {}", matchingLine, matchingMethod.name);
-      return matchingMethod;
+    List<MethodNode> matchingMethods = classFileLines.getMethodsByLine(matchingLine);
+    if (matchingMethods != null) {
+      matchingMethods.forEach(
+          methodNode -> {
+            log.debug("Found lineNode {} method: {}", matchingLine, methodNode.name);
+          });
+      // pick the first matching method.
+      // TODO need a way to disambiguate if multiple methods match the same line
+      return matchingMethods.isEmpty() ? null : matchingMethods.get(0);
     }
     log.debug("Cannot find line: {} in class {}", matchingLine, classNode.name);
     return null;

diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/Where.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/Where.java
@@ -64,8 +64,13 @@ protected static SourceLine[] sourceLines(String[] defs) {
 
   public static Where convertLineToMethod(Where lineWhere, ClassFileLines classFileLines) {
     if (lineWhere.methodName != null && lineWhere.lines != null) {
-      MethodNode method = classFileLines.getMethodByLine(lineWhere.lines[0].getFrom());
-      return new Where(lineWhere.typeName, method.name, method.desc, (SourceLine[]) null, null);
+      List<MethodNode> methodsByLine =
+          classFileLines.getMethodsByLine(lineWhere.lines[0].getFrom());
+      if (methodsByLine != null && !methodsByLine.isEmpty()) {
+        // pick the first method, as we can have multiple methods (lambdas) on the same line
+        MethodNode method = methodsByLine.get(0);
+        return new Where(lineWhere.typeName, method.name, method.desc, (SourceLine[]) null, null);
+      }
     }
     throw new IllegalArgumentException("Invalid where to convert from line to method " + lineWhere);
   }
@@ -122,7 +127,11 @@ public boolean isMethodMatching(MethodNode methodNode, ClassFileLines classFileL
         return true;
       }
       // try matching by line
-      return classFileLines.getMethodByLine(lines[0].getFrom()) == methodNode;
+      List<MethodNode> methodsByLine = classFileLines.getMethodsByLine(lines[0].getFrom());
+      if (methodsByLine == null || methodsByLine.isEmpty()) {
+        return false;
+      }
+      return methodsByLine.stream().anyMatch(m -> m == methodNode);
     }
     if (signature.equals("*") || signature.equals(targetMethodDescriptor)) {
       return true;

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
@@ -1,6 +1,8 @@
 package com.datadog.debugger.util;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -10,7 +12,7 @@
 import org.objectweb.asm.tree.MethodNode;
 
 public class ClassFileLines {
-  private final Map<Integer, MethodNode> methodByLine = new HashMap<>();
+  private final Map<Integer, List<MethodNode>> methodByLine = new HashMap<>();
   private final TreeMap<Integer, LabelNode> lineLabels = new TreeMap<>();
 
   public ClassFileLines(ClassNode classNode) {
@@ -19,7 +21,14 @@ public ClassFileLines(ClassNode classNode) {
       while (currentNode != null) {
         if (currentNode.getType() == AbstractInsnNode.LINE) {
           LineNumberNode lineNode = (LineNumberNode) currentNode;
-          methodByLine.put(lineNode.line, methodNode);
+          // on the same line, we can have multiple methods (lambdas, inner classes, etc)
+          List<MethodNode> methodNodes =
+              methodByLine.computeIfAbsent(lineNode.line, k -> new ArrayList<>());
+          // We are not using a Set here to keep the order of the methods
+          // We assume also that the number of methods per line is small
+          if (!methodNodes.contains(methodNode)) {
+            methodNodes.add(methodNode);
+          }
           // we are putting the line only the first time in the map.
           // for-loops can generate 2 line nodes for the same line, 1 before the loop, 1 for the
           // incrementation part.
@@ -31,7 +40,7 @@ public ClassFileLines(ClassNode classNode) {
     }
   }
 
-  public MethodNode getMethodByLine(int line) {
+  public List<MethodNode> getMethodsByLine(int line) {
     return methodByLine.get(line);
   }
 

diff --git a/...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
@@ -1443,12 +1443,12 @@ public void staticLambda() throws IOException, URISyntaxException {
     setCorrelationSingleton(spyCorrelationAccess);
     doReturn(true).when(spyCorrelationAccess).isAvailable();
     TestSnapshotListener listener =
-        installProbes(CLASS_NAME, createProbe(PROBE_ID, CLASS_NAME, null, null, "33"));
+        installProbes(CLASS_NAME, createProbe(PROBE_ID, CLASS_NAME, null, null, "37"));
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.onClass(testClass).call("main", "static", "email@address").get();
     assertEquals(8, result);
     Snapshot snapshot = assertOneSnapshot(listener);
-    CapturedContext context = snapshot.getCaptures().getLines().get(33);
+    CapturedContext context = snapshot.getCaptures().getLines().get(37);
     assertNotNull(context);
     assertCaptureLocals(context, "idx", "int", "5");
   }
@@ -1460,17 +1460,34 @@ public void capturingLambda() throws IOException, URISyntaxException {
     setCorrelationSingleton(spyCorrelationAccess);
     doReturn(true).when(spyCorrelationAccess).isAvailable();
     TestSnapshotListener listener =
-        installProbes(CLASS_NAME, createProbe(PROBE_ID, CLASS_NAME, null, null, "44"));
+        installProbes(CLASS_NAME, createProbe(PROBE_ID, CLASS_NAME, null, null, "48"));
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.onClass(testClass).call("main", "capturing", "email@address").get();
     assertEquals(8, result);
     Snapshot snapshot = assertOneSnapshot(listener);
-    CapturedContext context = snapshot.getCaptures().getLines().get(44);
+    CapturedContext context = snapshot.getCaptures().getLines().get(48);
     assertNotNull(context);
     assertCaptureLocals(context, "idx", "int", "5");
     assertCaptureFields(context, "strValue", "java.lang.String", "email@address");
   }
 
+  @Test
+  public void multiLambdas() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "CapturedSnapshot07";
+    CorrelationAccess spyCorrelationAccess = spy(CorrelationAccess.instance());
+    setCorrelationSingleton(spyCorrelationAccess);
+    doReturn(true).when(spyCorrelationAccess).isAvailable();
+    TestSnapshotListener listener =
+        installProbes(CLASS_NAME, createProbe(PROBE_ID, CLASS_NAME, null, null, "58"));
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "multi", "FOO1,FOO2,FOO3").get();
+    assertEquals(3, result);
+    Snapshot snapshot = assertOneSnapshot(listener);
+    CapturedContext context = snapshot.getCaptures().getLines().get(58);
+    assertNotNull(context);
+    assertCaptureArgs(context, "arg", String.class.getTypeName(), "FOO1,FOO2,FOO3");
+  }
+
   @Test
   public void tracerInstrumentedClass() throws Exception {
     DebuggerContext.initClassFilter(new DenyListHelper(null));

diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/WhereTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/WhereTest.java
@@ -8,6 +8,7 @@
 import com.datadog.debugger.util.MoshiHelper;
 import com.squareup.moshi.JsonAdapter;
 import java.io.IOException;
+import java.util.Arrays;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.objectweb.asm.tree.MethodNode;
@@ -119,7 +120,7 @@ public void methodMatching() {
     where = new Where("MyClass", "myMethod", null, new String[] {"42"}, null);
     ClassFileLines classFileLines = mock(ClassFileLines.class);
     MethodNode myMethodNode = createMethodNode("myMethod", "()V");
-    when(classFileLines.getMethodByLine(42)).thenReturn(myMethodNode);
+    when(classFileLines.getMethodsByLine(42)).thenReturn(Arrays.asList(myMethodNode));
     assertTrue(where.isMethodMatching(myMethodNode, classFileLines));
   }
 

diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/ClassFileLinesTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/ClassFileLinesTest.java
@@ -5,6 +5,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static utils.InstrumentationTestHelper.compile;
 
+import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.tree.ClassNode;
@@ -21,7 +22,7 @@ void isEmptyReturnTrue() {
 
   @Test
   void ValidateLineMapReturnLinesInRangeOrNull() throws Exception {
-    ClassFileLines classFileLines = createClassFileLines();
+    ClassFileLines classFileLines = createClassFileLines("CapturedSnapshot02");
     LabelNode line13 = classFileLines.getLineLabel(13);
     LabelNode line17 = classFileLines.getLineLabel(17);
 
@@ -31,34 +32,54 @@ void ValidateLineMapReturnLinesInRangeOrNull() throws Exception {
   }
 
   @Test
-  void getMethodByLine() throws Exception {
-    ClassFileLines classFileLines = createClassFileLines();
-    MethodNode method13 = classFileLines.getMethodByLine(13);
+  void getMethodsByLine() throws Exception {
+    ClassFileLines classFileLines = createClassFileLines("CapturedSnapshot02");
+    MethodNode method13 = classFileLines.getMethodsByLine(13).get(0);
     assertEquals("<init>", method13.name);
     assertEquals("()V", method13.desc);
-    MethodNode method17 = classFileLines.getMethodByLine(17);
+    MethodNode method17 = classFileLines.getMethodsByLine(17).get(0);
     assertEquals("<init>", method17.name);
     assertEquals("(Ljava/lang/Throwable;)V", method17.desc);
-    MethodNode method26 = classFileLines.getMethodByLine(26);
+    MethodNode method26 = classFileLines.getMethodsByLine(26).get(0);
     assertEquals("<init>", method26.name);
     assertEquals("(Ljava/lang/String;Ljava/lang/Object;)V", method26.desc);
-    MethodNode method32 = classFileLines.getMethodByLine(32);
+    MethodNode method32 = classFileLines.getMethodsByLine(32).get(0);
     assertEquals("createObject", method32.name);
     assertEquals("()Ljava/lang/Object;", method32.desc);
-    MethodNode method37 = classFileLines.getMethodByLine(37);
+    MethodNode method37 = classFileLines.getMethodsByLine(37).get(0);
     assertEquals("f", method37.name);
     assertEquals("()V", method37.desc);
-    MethodNode method44 = classFileLines.getMethodByLine(44);
+    MethodNode method44 = classFileLines.getMethodsByLine(44).get(0);
     assertEquals("synchronizedBlock", method44.name);
     assertEquals("(I)I", method44.desc);
-    MethodNode method54 = classFileLines.getMethodByLine(54);
+    MethodNode method54 = classFileLines.getMethodsByLine(54).get(0);
     assertEquals("main", method54.name);
     assertEquals("(Ljava/lang/String;)I", method54.desc);
   }
 
-  private ClassFileLines createClassFileLines() throws Exception {
-    final String CLASS_NAME = "CapturedSnapshot02";
-    byte[] byteBuffer = compile(CLASS_NAME).get(CLASS_NAME);
+  @Test
+  void getMethodsByLineWithLambdas() throws Exception {
+    ClassFileLines classFileLines = createClassFileLines("CapturedSnapshot07");
+    List<MethodNode> methods57 = classFileLines.getMethodsByLine(57);
+    // despite having the return on multi line which yield multiple BCI for same line
+    // this is still the same method
+    assertEquals(1, methods57.size());
+    assertEquals("multiLambda", methods57.get(0).name);
+    assertEquals("(Ljava/lang/String;)I", methods57.get(0).desc);
+    List<MethodNode> methods58 = classFileLines.getMethodsByLine(58);
+    assertEquals(3, methods58.size());
+    assertEquals("multiLambda", methods58.get(0).name);
+    assertEquals("(Ljava/lang/String;)I", methods58.get(0).desc);
+    // filter
+    assertEquals("lambda$multiLambda$3", methods58.get(1).name);
+    assertEquals("(Ljava/lang/String;)Z", methods58.get(1).desc);
+    // map
+    assertEquals("lambda$multiLambda$2", methods58.get(2).name);
+    assertEquals("(Ljava/lang/String;)Ljava/lang/String;", methods58.get(2).desc);
+  }
+
+  private ClassFileLines createClassFileLines(String className) throws Exception {
+    byte[] byteBuffer = compile(className).get(className);
     ClassReader reader = new ClassReader(byteBuffer);
     ClassNode classNode = new ClassNode();
     reader.accept(classNode, ClassReader.SKIP_FRAMES);

diff --git a/dd-java-agent/agent-debugger/src/test/resources/CapturedSnapshot07.java b/dd-java-agent/agent-debugger/src/test/resources/CapturedSnapshot07.java
@@ -24,6 +24,10 @@ public static int main(String mode, String arg) {
       cs7.strValue = arg;
       return cs7.capturingLambda(arg);
     }
+    if (mode.equals("multi")) {
+      CapturedSnapshot07 cs7 = new CapturedSnapshot07();
+      return cs7.multiLambda(arg);
+    }
     return 0;
   }
 
@@ -48,4 +52,10 @@ private int capturingLambda(String arg) {
     };
     return func.apply(arg).length();
   }
+
+  private int multiLambda(String arg) {
+    return (int)Arrays.stream(arg.split(","))
+        .map(s -> s.toUpperCase()).filter(s -> s.startsWith("FOO"))
+        .count();
+  }
 }