Refactor probe matching for methods (#8021)

To be able to match probe with more flexibility (probes targeting a
method with or without signature) we are iterating over all methods
of a class and using Where::isMethodMatching to select probe
definitions that will be apply on the current method.
With this new approach we are dropping the support of overloaded
methods: only one method is instrumented or a definition is applied
only once.

Author: jpbempel

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java

@@ -240,22 +240,6 @@ public byte[] transform(
     return null;
   }
 
-  private Map<Where, List<ProbeDefinition>> mergeLocations(
-      List<ProbeDefinition> definitions, ClassFileLines classFileLines) {
-    Map<Where, List<ProbeDefinition>> mergedProbes = new HashMap<>();
-    for (ProbeDefinition definition : definitions) {
-      Where where = definition.getWhere();
-      if (definition instanceof ForceMethodInstrumentation) {
-        // normalize where for line => to precise method location
-        where = Where.convertLineToMethod(definition.getWhere(), classFileLines);
-      }
-      List<ProbeDefinition> instrumentationDefinitions =
-          mergedProbes.computeIfAbsent(where, key -> new ArrayList<>());
-      instrumentationDefinitions.add(definition);
-    }
-    return mergedProbes;
-  }
-
   private boolean skipInstrumentation(ClassLoader loader, String classFilePath) {
     if (definitionMatcher.isEmpty()) {
       log.warn("No debugger definitions present.");
@@ -496,41 +480,35 @@ private boolean performInstrumentation(
       ClassNode classNode) {
     boolean transformed = false;
     ClassFileLines classFileLines = new ClassFileLines(classNode);
-    Map<Where, List<ProbeDefinition>> definitionByLocation =
-        mergeLocations(definitions, classFileLines);
-    // FIXME build a map also for methods to optimize the matching, currently O(probes*methods)
-    for (Map.Entry<Where, List<ProbeDefinition>> entry : definitionByLocation.entrySet()) {
-      Where where = entry.getKey();
-      String methodName = where.getMethodName();
-      String[] lines = where.getLines();
-      List<MethodNode> methodNodes;
-      if (methodName == null && lines != null) {
-        MethodNode methodNode = matchSourceFile(classNode, where, classFileLines);
-        methodNodes = methodNode != null ? singletonList(methodNode) : Collections.emptyList();
-      } else {
-        methodNodes = matchMethodDescription(classNode, where, classFileLines);
+    Set<ProbeDefinition> remainingDefinitions = new HashSet<>(definitions);
+    for (MethodNode methodNode : classNode.methods) {
+      List<ProbeDefinition> matchingDefs = new ArrayList<>();
+      for (ProbeDefinition definition : definitions) {
+        if (definition.getWhere().isMethodMatching(methodNode, classFileLines)
+            && remainingDefinitions.contains(definition)) {
+          matchingDefs.add(definition);
+          remainingDefinitions.remove(definition);
+        }
       }
-      List<ProbeDefinition> definitionsByWhere = entry.getValue();
-      if (methodNodes.isEmpty()) {
-        reportLocationNotFound(definitionsByWhere, classNode.name, methodName);
+      if (matchingDefs.isEmpty()) {
         continue;
       }
-      for (MethodNode methodNode : methodNodes) {
-        if (log.isDebugEnabled()) {
-          List<String> probeIds =
-              definitionsByWhere.stream().map(ProbeDefinition::getId).collect(toList());
-          log.debug(
-              "Instrumenting method: {}.{}{} for probe ids: {}",
-              fullyQualifiedClassName,
-              methodNode.name,
-              methodNode.desc,
-              probeIds);
-        }
-        MethodInfo methodInfo = new MethodInfo(loader, classNode, methodNode, classFileLines);
-        InstrumentationResult result = applyInstrumentation(methodInfo, definitionsByWhere);
-        transformed |= result.isInstalled();
-        handleInstrumentationResult(definitionsByWhere, result);
+      if (log.isDebugEnabled()) {
+        List<String> probeIds = matchingDefs.stream().map(ProbeDefinition::getId).collect(toList());
+        log.debug(
+            "Instrumenting method: {}.{}{} for probe ids: {}",
+            fullyQualifiedClassName,
+            methodNode.name,
+            methodNode.desc,
+            probeIds);
       }
+      MethodInfo methodInfo = new MethodInfo(loader, classNode, methodNode, classFileLines);
+      InstrumentationResult result = applyInstrumentation(methodInfo, matchingDefs);
+      transformed |= result.isInstalled();
+      handleInstrumentationResult(matchingDefs, result);
+    }
+    for (ProbeDefinition definition : remainingDefinitions) {
+      reportLocationNotFound(definition, classNode.name, definition.getWhere().getMethodName());
     }
     return transformed;
   }
@@ -556,9 +534,9 @@ private void handleInstrumentationResult(
   }
 
   private void reportLocationNotFound(
-      List<ProbeDefinition> definitions, String className, String methodName) {
+      ProbeDefinition definition, String className, String methodName) {
     if (methodName != null) {
-      reportErrorForAllProbes(definitions, CANNOT_FIND_METHOD, className, methodName);
+      reportErrorForAllProbes(singletonList(definition), CANNOT_FIND_METHOD, className, methodName);
       return;
     }
     // This is a line probe, so we don't report line not found because the line may be found later
@@ -628,39 +606,44 @@ static class ToInstrumentInfo {
     }
   }
 
+  private static boolean isCapturedContextProbe(ProbeDefinition definition) {
+    return definition instanceof LogProbe
+        || definition instanceof SpanDecorationProbe
+        || definition instanceof TriggerProbe;
+  }
+
   private List<ToInstrumentInfo> filterAndSortDefinitions(
       List<ProbeDefinition> definitions, ClassFileLines classFileLines) {
     List<ToInstrumentInfo> toInstrument = new ArrayList<>();
     List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
+    Map<Where, List<ProbeDefinition>> capturedContextLineProbes = new HashMap<>();
     boolean addedExceptionProbe = false;
     for (ProbeDefinition definition : definitions) {
       // Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
       // and therefore need to be instrumented once
       // note: exception probes are log probes and are handled the same way
-      if (definition instanceof LogProbe
-          || definition instanceof SpanDecorationProbe
-          || definition instanceof TriggerProbe) {
-        if (definition instanceof ExceptionProbe) {
-          if (addedExceptionProbe) {
-            continue;
+      if (isCapturedContextProbe(definition)) {
+        if (definition.isLineProbe()) {
+          capturedContextLineProbes
+              .computeIfAbsent(definition.getWhere(), key -> new ArrayList<>())
+              .add(definition);
+        } else {
+          if (definition instanceof ExceptionProbe) {
+            if (addedExceptionProbe) {
+              continue;
+            }
+            // only add one exception probe to the list of probes to instrument
+            // to have only one instance (1 probeId) of exception probe to handle all exceptions
+            addedExceptionProbe = true;
           }
-          // only add one exception probe to the list of probes to instrument
-          // to have only one instance (1 probeId) of exception probe to handle all exceptions
-          addedExceptionProbe = true;
+          capturedContextProbes.add(definition);
         }
-        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, classFileLines);
-      toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
-    }
-    //     LOGGER.debug("exception probe is already instrumented for {}", preciseWhere);
+    processCapturedContextLineProbes(capturedContextLineProbes, toInstrument);
+    processCapturedContextMethodProbes(classFileLines, capturedContextProbes, toInstrument);
     // ordering: metric < log < span decoration < span
     toInstrument.sort(
         (info1, info2) -> {
@@ -671,6 +654,32 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(
     return toInstrument;
   }
 
+  private void processCapturedContextMethodProbes(
+      ClassFileLines classFileLines,
+      List<ProbeDefinition> capturedContextProbes,
+      List<ToInstrumentInfo> toInstrument) {
+    if (capturedContextProbes.isEmpty()) {
+      return;
+    }
+    List<ProbeId> probesIds =
+        capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
+    ProbeDefinition referenceDefinition =
+        selectReferenceDefinition(capturedContextProbes, classFileLines);
+    toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
+  }
+
+  private static void processCapturedContextLineProbes(
+      Map<Where, List<ProbeDefinition>> lineProbes, List<ToInstrumentInfo> toInstrument) {
+    for (Map.Entry<Where, List<ProbeDefinition>> entry : lineProbes.entrySet()) {
+      if (entry.getValue().isEmpty()) {
+        continue;
+      }
+      List<ProbeId> probeIds =
+          entry.getValue().stream().map(ProbeDefinition::getProbeId).collect(toList());
+      toInstrument.add(new ToInstrumentInfo(entry.getValue().get(0), probeIds));
+    }
+  }
+
   // Log & Span Decoration probes share the same instrumentor so only one definition should be
   // used to generate the instrumentation. This method generate a synthetic definition that
   // match the type of the probe to instrument: if at least one probe is LogProbe then we are

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java

@@ -1360,6 +1360,22 @@ public void mergedProbesConditionMixedLocation() throws IOException, URISyntaxEx
     assertNull(listener.snapshots.get(1).getEvaluationErrors());
   }
 
+  @Test
+  public void mergedProbesDifferentSignature() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "CapturedSnapshot08";
+    LogProbe probe1 = createProbeAtExit(PROBE_ID1, CLASS_NAME, "doit", null);
+    LogProbe probe2 = createProbeAtExit(PROBE_ID2, CLASS_NAME, "doit", "int (java.lang.String)");
+    LogProbe probe3 = createProbeAtExit(PROBE_ID3, CLASS_NAME, "doit", "(String)");
+    TestSnapshotListener listener = installProbes(probe1, probe2, probe3);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "1").get();
+    Assertions.assertEquals(3, result);
+    Assertions.assertEquals(3, listener.snapshots.size());
+    assertNull(listener.snapshots.get(0).getEvaluationErrors());
+    assertNull(listener.snapshots.get(1).getEvaluationErrors());
+    assertNull(listener.snapshots.get(2).getEvaluationErrors());
+  }
+
   @Test
   public void fields() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot06";
@@ -1579,11 +1595,9 @@ public void overloadedMethods() throws Exception {
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.onClass(testClass).call("main", "").get();
     assertEquals(63, result);
-    List<Snapshot> snapshots = assertSnapshots(listener, 4, PROBE_ID, PROBE_ID, PROBE_ID, PROBE_ID);
+    List<Snapshot> snapshots = assertSnapshots(listener, 1, PROBE_ID);
     assertCaptureReturnValue(snapshots.get(0).getCaptures().getReturn(), "int", "42");
-    assertCaptureArgs(snapshots.get(1).getCaptures().getEntry(), "s", "java.lang.String", "1");
-    assertCaptureArgs(snapshots.get(2).getCaptures().getEntry(), "s", "java.lang.String", "2");
-    assertCaptureArgs(snapshots.get(3).getCaptures().getEntry(), "s", "java.lang.String", "3");
+    assertEquals(1, snapshots.get(0).getCaptures().getEntry().getArguments().size());
   }
 
   @Test