Skip to content

Commit f2cc473

Browse files
jpbempeljpbempel
committed
Refactor probe matching for methods
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.
1 parent 565afcd commit f2cc473

File tree

2 files changed

+94
-71
lines changed

2 files changed

+94
-71
lines changed

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

+76-67
Original file line numberDiff line numberDiff line change
@@ -240,22 +240,6 @@ public byte[] transform(
240240
return null;
241241
}
242242

243-
private Map<Where, List<ProbeDefinition>> mergeLocations(
244-
List<ProbeDefinition> definitions, ClassFileLines classFileLines) {
245-
Map<Where, List<ProbeDefinition>> mergedProbes = new HashMap<>();
246-
for (ProbeDefinition definition : definitions) {
247-
Where where = definition.getWhere();
248-
if (definition instanceof ForceMethodInstrumentation) {
249-
// normalize where for line => to precise method location
250-
where = Where.convertLineToMethod(definition.getWhere(), classFileLines);
251-
}
252-
List<ProbeDefinition> instrumentationDefinitions =
253-
mergedProbes.computeIfAbsent(where, key -> new ArrayList<>());
254-
instrumentationDefinitions.add(definition);
255-
}
256-
return mergedProbes;
257-
}
258-
259243
private boolean skipInstrumentation(ClassLoader loader, String classFilePath) {
260244
if (definitionMatcher.isEmpty()) {
261245
log.warn("No debugger definitions present.");
@@ -496,41 +480,35 @@ private boolean performInstrumentation(
496480
ClassNode classNode) {
497481
boolean transformed = false;
498482
ClassFileLines classFileLines = new ClassFileLines(classNode);
499-
Map<Where, List<ProbeDefinition>> definitionByLocation =
500-
mergeLocations(definitions, classFileLines);
501-
// FIXME build a map also for methods to optimize the matching, currently O(probes*methods)
502-
for (Map.Entry<Where, List<ProbeDefinition>> entry : definitionByLocation.entrySet()) {
503-
Where where = entry.getKey();
504-
String methodName = where.getMethodName();
505-
String[] lines = where.getLines();
506-
List<MethodNode> methodNodes;
507-
if (methodName == null && lines != null) {
508-
MethodNode methodNode = matchSourceFile(classNode, where, classFileLines);
509-
methodNodes = methodNode != null ? singletonList(methodNode) : Collections.emptyList();
510-
} else {
511-
methodNodes = matchMethodDescription(classNode, where, classFileLines);
483+
Set<ProbeDefinition> remainingDefinitions = new HashSet<>(definitions);
484+
for (MethodNode methodNode : classNode.methods) {
485+
List<ProbeDefinition> matchingDefs = new ArrayList<>();
486+
for (ProbeDefinition definition : definitions) {
487+
if (definition.getWhere().isMethodMatching(methodNode, classFileLines)
488+
&& remainingDefinitions.contains(definition)) {
489+
matchingDefs.add(definition);
490+
remainingDefinitions.remove(definition);
491+
}
512492
}
513-
List<ProbeDefinition> definitionsByWhere = entry.getValue();
514-
if (methodNodes.isEmpty()) {
515-
reportLocationNotFound(definitionsByWhere, classNode.name, methodName);
493+
if (matchingDefs.isEmpty()) {
516494
continue;
517495
}
518-
for (MethodNode methodNode : methodNodes) {
519-
if (log.isDebugEnabled()) {
520-
List<String> probeIds =
521-
definitionsByWhere.stream().map(ProbeDefinition::getId).collect(toList());
522-
log.debug(
523-
"Instrumenting method: {}.{}{} for probe ids: {}",
524-
fullyQualifiedClassName,
525-
methodNode.name,
526-
methodNode.desc,
527-
probeIds);
528-
}
529-
MethodInfo methodInfo = new MethodInfo(loader, classNode, methodNode, classFileLines);
530-
InstrumentationResult result = applyInstrumentation(methodInfo, definitionsByWhere);
531-
transformed |= result.isInstalled();
532-
handleInstrumentationResult(definitionsByWhere, result);
496+
if (log.isDebugEnabled()) {
497+
List<String> probeIds = matchingDefs.stream().map(ProbeDefinition::getId).collect(toList());
498+
log.debug(
499+
"Instrumenting method: {}.{}{} for probe ids: {}",
500+
fullyQualifiedClassName,
501+
methodNode.name,
502+
methodNode.desc,
503+
probeIds);
533504
}
505+
MethodInfo methodInfo = new MethodInfo(loader, classNode, methodNode, classFileLines);
506+
InstrumentationResult result = applyInstrumentation(methodInfo, matchingDefs);
507+
transformed |= result.isInstalled();
508+
handleInstrumentationResult(matchingDefs, result);
509+
}
510+
for (ProbeDefinition definition : remainingDefinitions) {
511+
reportLocationNotFound(definition, classNode.name, definition.getWhere().getMethodName());
534512
}
535513
return transformed;
536514
}
@@ -556,9 +534,9 @@ private void handleInstrumentationResult(
556534
}
557535

558536
private void reportLocationNotFound(
559-
List<ProbeDefinition> definitions, String className, String methodName) {
537+
ProbeDefinition definition, String className, String methodName) {
560538
if (methodName != null) {
561-
reportErrorForAllProbes(definitions, CANNOT_FIND_METHOD, className, methodName);
539+
reportErrorForAllProbes(singletonList(definition), CANNOT_FIND_METHOD, className, methodName);
562540
return;
563541
}
564542
// 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 {
628606
}
629607
}
630608

609+
private static boolean isCapturedContextProbe(ProbeDefinition definition) {
610+
return definition instanceof LogProbe
611+
|| definition instanceof SpanDecorationProbe
612+
|| definition instanceof TriggerProbe;
613+
}
614+
631615
private List<ToInstrumentInfo> filterAndSortDefinitions(
632616
List<ProbeDefinition> definitions, ClassFileLines classFileLines) {
633617
List<ToInstrumentInfo> toInstrument = new ArrayList<>();
634618
List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
619+
Map<Where, List<ProbeDefinition>> capturedContextLineProbes = new HashMap<>();
635620
boolean addedExceptionProbe = false;
636621
for (ProbeDefinition definition : definitions) {
637622
// Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
638623
// and therefore need to be instrumented once
639624
// note: exception probes are log probes and are handled the same way
640-
if (definition instanceof LogProbe
641-
|| definition instanceof SpanDecorationProbe
642-
|| definition instanceof TriggerProbe) {
643-
if (definition instanceof ExceptionProbe) {
644-
if (addedExceptionProbe) {
645-
continue;
625+
if (isCapturedContextProbe(definition)) {
626+
if (definition.isLineProbe()) {
627+
capturedContextLineProbes
628+
.computeIfAbsent(definition.getWhere(), key -> new ArrayList<>())
629+
.add(definition);
630+
} else {
631+
if (definition instanceof ExceptionProbe) {
632+
if (addedExceptionProbe) {
633+
continue;
634+
}
635+
// only add one exception probe to the list of probes to instrument
636+
// to have only one instance (1 probeId) of exception probe to handle all exceptions
637+
addedExceptionProbe = true;
646638
}
647-
// only add one exception probe to the list of probes to instrument
648-
// to have only one instance (1 probeId) of exception probe to handle all exceptions
649-
addedExceptionProbe = true;
639+
capturedContextProbes.add(definition);
650640
}
651-
capturedContextProbes.add(definition);
652641
} else {
653642
toInstrument.add(new ToInstrumentInfo(definition, singletonList(definition.getProbeId())));
654643
}
655644
}
656-
if (!capturedContextProbes.isEmpty()) {
657-
List<ProbeId> probesIds =
658-
capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
659-
ProbeDefinition referenceDefinition =
660-
selectReferenceDefinition(capturedContextProbes, classFileLines);
661-
toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
662-
}
663-
// LOGGER.debug("exception probe is already instrumented for {}", preciseWhere);
645+
processCapturedContextLineProbes(capturedContextLineProbes, toInstrument);
646+
processCapturedContextMethodProbes(classFileLines, capturedContextProbes, toInstrument);
664647
// ordering: metric < log < span decoration < span
665648
toInstrument.sort(
666649
(info1, info2) -> {
@@ -671,6 +654,32 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(
671654
return toInstrument;
672655
}
673656

657+
private void processCapturedContextMethodProbes(
658+
ClassFileLines classFileLines,
659+
List<ProbeDefinition> capturedContextProbes,
660+
List<ToInstrumentInfo> toInstrument) {
661+
if (capturedContextProbes.isEmpty()) {
662+
return;
663+
}
664+
List<ProbeId> probesIds =
665+
capturedContextProbes.stream().map(ProbeDefinition::getProbeId).collect(toList());
666+
ProbeDefinition referenceDefinition =
667+
selectReferenceDefinition(capturedContextProbes, classFileLines);
668+
toInstrument.add(new ToInstrumentInfo(referenceDefinition, probesIds));
669+
}
670+
671+
private static void processCapturedContextLineProbes(
672+
Map<Where, List<ProbeDefinition>> lineProbes, List<ToInstrumentInfo> toInstrument) {
673+
for (Map.Entry<Where, List<ProbeDefinition>> entry : lineProbes.entrySet()) {
674+
if (entry.getValue().isEmpty()) {
675+
continue;
676+
}
677+
List<ProbeId> probeIds =
678+
entry.getValue().stream().map(ProbeDefinition::getProbeId).collect(toList());
679+
toInstrument.add(new ToInstrumentInfo(entry.getValue().get(0), probeIds));
680+
}
681+
}
682+
674683
// Log & Span Decoration probes share the same instrumentor so only one definition should be
675684
// used to generate the instrumentation. This method generate a synthetic definition that
676685
// 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

+18-4
Original file line numberDiff line numberDiff line change
@@ -1360,6 +1360,22 @@ public void mergedProbesConditionMixedLocation() throws IOException, URISyntaxEx
13601360
assertNull(listener.snapshots.get(1).getEvaluationErrors());
13611361
}
13621362

1363+
@Test
1364+
public void mergedProbesDifferentSignature() throws IOException, URISyntaxException {
1365+
final String CLASS_NAME = "CapturedSnapshot08";
1366+
LogProbe probe1 = createProbeAtExit(PROBE_ID1, CLASS_NAME, "doit", null);
1367+
LogProbe probe2 = createProbeAtExit(PROBE_ID2, CLASS_NAME, "doit", "int (java.lang.String)");
1368+
LogProbe probe3 = createProbeAtExit(PROBE_ID3, CLASS_NAME, "doit", "(String)");
1369+
TestSnapshotListener listener = installProbes(probe1, probe2, probe3);
1370+
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
1371+
int result = Reflect.onClass(testClass).call("main", "1").get();
1372+
Assertions.assertEquals(3, result);
1373+
Assertions.assertEquals(3, listener.snapshots.size());
1374+
assertNull(listener.snapshots.get(0).getEvaluationErrors());
1375+
assertNull(listener.snapshots.get(1).getEvaluationErrors());
1376+
assertNull(listener.snapshots.get(2).getEvaluationErrors());
1377+
}
1378+
13631379
@Test
13641380
public void fields() throws IOException, URISyntaxException {
13651381
final String CLASS_NAME = "CapturedSnapshot06";
@@ -1579,11 +1595,9 @@ public void overloadedMethods() throws Exception {
15791595
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
15801596
int result = Reflect.onClass(testClass).call("main", "").get();
15811597
assertEquals(63, result);
1582-
List<Snapshot> snapshots = assertSnapshots(listener, 4, PROBE_ID, PROBE_ID, PROBE_ID, PROBE_ID);
1598+
List<Snapshot> snapshots = assertSnapshots(listener, 1, PROBE_ID);
15831599
assertCaptureReturnValue(snapshots.get(0).getCaptures().getReturn(), "int", "42");
1584-
assertCaptureArgs(snapshots.get(1).getCaptures().getEntry(), "s", "java.lang.String", "1");
1585-
assertCaptureArgs(snapshots.get(2).getCaptures().getEntry(), "s", "java.lang.String", "2");
1586-
assertCaptureArgs(snapshots.get(3).getCaptures().getEntry(), "s", "java.lang.String", "3");
1600+
assertEquals(1, snapshots.get(0).getCaptures().getEntry().getArguments().size());
15871601
}
15881602

15891603
@Test

0 commit comments

Comments
 (0)