Skip to content

Commit

Permalink
Refactor probe matching for methods (#8021)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@jpbempel @PerfectSlayer
jpbempel authored and PerfectSlayer committed Nov 28, 2024
1 parent 8c921d8 commit 4e85847
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 71 deletions.
143 changes: 76 additions & 67 deletions ...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand Down Expand Up @@ -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;
}
Expand All @@ -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
Expand Down Expand Up @@ -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) -> {
Expand All @@ -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
Expand Down
22 changes: 18 additions & 4 deletions ...a-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 4e85847

Please sign in to comment.