Skip to content

Commit 617b47a

Browse files
jpbempeljpbempel
committed
Fix Capture Expressions support for multi-probes
When multiple probes are created on the same location we previously assume that the context can be shared for all snapshots of the probes. With Capture Expressions we now need to differentiate the captures to respect the probe definitions. We are now creating CapturedContext specifically based on the probe definiton and we are filtering Capture Expressions.
1 parent 9c3d035 commit 617b47a

4 files changed

Lines changed: 179 additions & 2 deletions

File tree

‎dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/CapturedContext.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,11 @@ private CapturedContext(CapturedContext other, Map<String, Object> extensions) {
5353
this.arguments = other.arguments;
5454
this.locals = other.getLocals();
5555
this.throwable = other.throwable;
56+
this.staticFields = other.staticFields;
57+
this.limits = other.limits;
58+
this.thisClassName = other.thisClassName;
59+
this.duration = other.duration;
60+
this.captureExpressions = other.captureExpressions;
5661
this.extensions.putAll(other.extensions);
5762
this.extensions.putAll(extensions);
5863
}
@@ -177,6 +182,14 @@ private Object tryRetrieve(String name) {
177182
return result != null ? result : Values.UNDEFINED_OBJECT;
178183
}
179184

185+
public CapturedContext copyWithoutCaptureExpressions() {
186+
CapturedContext newContext = new CapturedContext(this, Collections.emptyMap());
187+
if (newContext.captureExpressions != null) {
188+
newContext.captureExpressions = null;
189+
}
190+
return newContext;
191+
}
192+
180193
@Override
181194
public ValueReferenceResolver withExtensions(Map<String, Object> extensions) {
182195
return new CapturedContext(this, extensions);

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -781,6 +781,7 @@ private static ProbeDefinition selectReferenceDefinition(
781781
LogProbe.Capture capture = null;
782782
boolean captureSnapshot = false;
783783
ProbeCondition probeCondition = null;
784+
List<LogProbe.CaptureExpression> captureExpressions = null;
784785
Where where = capturedContextProbes.get(0).getWhere();
785786
ProbeId probeId = capturedContextProbes.get(0).getProbeId();
786787
for (ProbeDefinition definition : capturedContextProbes) {
@@ -792,6 +793,8 @@ private static ProbeDefinition selectReferenceDefinition(
792793
LogProbe logProbe = (LogProbe) definition;
793794
captureSnapshot = captureSnapshot | logProbe.isCaptureSnapshot();
794795
capture = mergeCapture(capture, logProbe.getCapture());
796+
// captureExpressions = mergeCaptureExpressions(captureExpressions,
797+
// logProbe.getCaptureExpressions());
795798
if (probeCondition == null) {
796799
probeCondition = logProbe.getProbeCondition();
797800
}
@@ -837,6 +840,19 @@ private static LogProbe.Capture mergeCapture(
837840
Math.max(current.getMaxFieldCount(), newCapture.getMaxFieldCount()));
838841
}
839842

843+
private static List<LogProbe.CaptureExpression> mergeCaptureExpressions(
844+
List<LogProbe.CaptureExpression> captureExpressions,
845+
List<LogProbe.CaptureExpression> newCaptureExpressions) {
846+
if (captureExpressions == null) {
847+
return newCaptureExpressions;
848+
}
849+
if (newCaptureExpressions == null) {
850+
return captureExpressions;
851+
}
852+
captureExpressions.addAll(newCaptureExpressions);
853+
return captureExpressions;
854+
}
855+
840856
private InstrumentationResult.Status preCheckInstrumentation(
841857
Map<ProbeId, List<DiagnosticMessage>> diagnostics, MethodInfo methodInfo) {
842858
if ((methodInfo.getMethodNode().access & (Opcodes.ACC_NATIVE | Opcodes.ACC_ABSTRACT)) != 0) {

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

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -636,8 +636,7 @@ protected boolean fillSnapshot(
636636
snapshot.setTraceId(CorrelationIdentifier.getTraceId());
637637
snapshot.setSpanId(CorrelationIdentifier.getSpanId());
638638
if (isFullSnapshot()) {
639-
snapshot.setEntry(entryContext);
640-
snapshot.setExit(exitContext);
639+
assignCaptures(snapshot, entryContext, exitContext);
641640
}
642641
snapshot.setMessage(message);
643642
snapshot.setDuration(exitContext.getDuration());
@@ -655,6 +654,41 @@ protected boolean fillSnapshot(
655654
return shouldCommit;
656655
}
657656

657+
private void assignCaptures(
658+
Snapshot snapshot, CapturedContext entryContext, CapturedContext exitContext) {
659+
if (isCaptureSnapshot()) {
660+
addContextWithoutCaptureExpressions(entryContext, snapshot::setEntry);
661+
addContextWithoutCaptureExpressions(exitContext, snapshot::setExit);
662+
} else if (captureExpressions != null) {
663+
addFilteredCaptureExpressions(entryContext, snapshot::setEntry);
664+
addFilteredCaptureExpressions(exitContext, snapshot::setExit);
665+
}
666+
}
667+
668+
private void addContextWithoutCaptureExpressions(
669+
CapturedContext context, Consumer<CapturedContext> setContext) {
670+
// no capture expressions, assign directly the context in the snapshot
671+
if (context.getCaptureExpressions() == null || context.getCaptureExpressions().isEmpty()) {
672+
setContext.accept(context);
673+
return;
674+
}
675+
CapturedContext newContext = context.copyWithoutCaptureExpressions();
676+
setContext.accept(newContext);
677+
}
678+
679+
private void addFilteredCaptureExpressions(
680+
CapturedContext capturedContext, Consumer<CapturedContext> setContext) {
681+
Map<String, CapturedContext.CapturedValue> contextCapExpr =
682+
capturedContext.getCaptureExpressions();
683+
if (contextCapExpr != null && !contextCapExpr.isEmpty()) {
684+
CapturedContext newContext = new CapturedContext();
685+
for (CaptureExpression capExprDef : captureExpressions) {
686+
newContext.addCaptureExpression(contextCapExpr.get(capExprDef.getName()));
687+
}
688+
setContext.accept(newContext);
689+
}
690+
}
691+
658692
private void processCaptureExpressions(CapturedContext context, LogStatus logStatus) {
659693
if (captureExpressions == null) {
660694
return;

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

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1522,6 +1522,120 @@ public void mergedProbesDifferentSignature() throws IOException, URISyntaxExcept
15221522
assertNull(listener.snapshots.get(2).getEvaluationErrors());
15231523
}
15241524

1525+
@Test
1526+
public void mergedProbesWithCaptureExpressionsMixed() throws IOException, URISyntaxException {
1527+
final String CLASS_NAME = "CapturedSnapshot08";
1528+
LogProbe probe1 =
1529+
createProbeBuilder(PROBE_ID1, CLASS_NAME, "doit", null)
1530+
.evaluateAt(MethodLocation.EXIT)
1531+
.captureSnapshot(false)
1532+
.captureExpressions(
1533+
Arrays.asList(
1534+
new LogProbe.CaptureExpression(
1535+
"typed_fld_fld_msg",
1536+
new ValueScript(
1537+
DSL.getMember(
1538+
DSL.getMember(DSL.getMember(DSL.ref("typed"), "fld"), "fld"),
1539+
"msg"),
1540+
"typed.fld.fld.msg"),
1541+
null),
1542+
new LogProbe.CaptureExpression(
1543+
"nullTyped_fld",
1544+
new ValueScript(
1545+
DSL.getMember(DSL.ref("nullTyped"), "fld"), "nullTyped.fld"),
1546+
null)))
1547+
.build();
1548+
LogProbe probe2 = createMethodProbeAtExit(PROBE_ID2, CLASS_NAME, "doit", null);
1549+
TestSnapshotListener listener = installProbes(probe1, probe2);
1550+
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
1551+
int result = Reflect.onClass(testClass).call("main", "1").get();
1552+
assertEquals(3, result);
1553+
List<Snapshot> snapshots = assertSnapshots(listener, 2, PROBE_ID1, PROBE_ID2);
1554+
// snapshot with Capture Expressions
1555+
Snapshot snapshot0 = snapshots.get(0);
1556+
assertEquals(2, snapshot0.getCaptures().getReturn().getCaptureExpressions().size());
1557+
assertCaptureExpressions(
1558+
snapshot0.getCaptures().getReturn(),
1559+
"typed_fld_fld_msg",
1560+
String.class.getTypeName(),
1561+
"hello");
1562+
assertCaptureExpressions(
1563+
snapshot0.getCaptures().getReturn(), "nullTyped_fld", Object.class.getTypeName(), null);
1564+
assertNull(snapshot0.getCaptures().getReturn().getArguments());
1565+
assertNull(snapshot0.getCaptures().getReturn().getLocals());
1566+
// Snapshot without Capture Expressions
1567+
Snapshot snapshot1 = snapshots.get(1);
1568+
assertNull(snapshot1.getCaptures().getReturn().getCaptureExpressions());
1569+
assertCaptureArgs(snapshot1.getCaptures().getReturn(), "arg", String.class.getTypeName(), "1");
1570+
assertCaptureLocals(
1571+
snapshot1.getCaptures().getReturn(), "var1", Integer.TYPE.getTypeName(), "3");
1572+
assertCaptureLocals(
1573+
snapshot1.getCaptures().getReturn(), "@return", Integer.TYPE.getTypeName(), "3");
1574+
}
1575+
1576+
@Test
1577+
public void mergedProbesWithDifferentCaptureExpressions() throws IOException, URISyntaxException {
1578+
final String CLASS_NAME = "CapturedSnapshot08";
1579+
LogProbe probe1 =
1580+
createProbeBuilder(PROBE_ID1, CLASS_NAME, "doit", null)
1581+
.evaluateAt(MethodLocation.EXIT)
1582+
.captureSnapshot(false)
1583+
.captureExpressions(
1584+
Arrays.asList(
1585+
new LogProbe.CaptureExpression(
1586+
"typed_fld_fld_msg",
1587+
new ValueScript(
1588+
DSL.getMember(
1589+
DSL.getMember(DSL.getMember(DSL.ref("typed"), "fld"), "fld"),
1590+
"msg"),
1591+
"typed.fld.fld.msg"),
1592+
null),
1593+
new LogProbe.CaptureExpression(
1594+
"nullTyped_fld",
1595+
new ValueScript(
1596+
DSL.getMember(DSL.ref("nullTyped"), "fld"), "nullTyped.fld"),
1597+
null)))
1598+
.build();
1599+
LogProbe probe2 =
1600+
createProbeBuilder(PROBE_ID2, CLASS_NAME, "doit", null)
1601+
.evaluateAt(MethodLocation.EXIT)
1602+
.captureSnapshot(false)
1603+
.captureExpressions(
1604+
Arrays.asList(
1605+
new LogProbe.CaptureExpression(
1606+
"var1", new ValueScript(DSL.ref("var1"), "var1"), null),
1607+
new LogProbe.CaptureExpression(
1608+
"this_fld",
1609+
new ValueScript(DSL.getMember(DSL.ref("this"), "fld"), "this.fld"),
1610+
null)))
1611+
.build();
1612+
TestSnapshotListener listener = installProbes(probe1, probe2);
1613+
Class<?> testClass = compileAndLoadClass(CLASS_NAME);
1614+
int result = Reflect.onClass(testClass).call("main", "1").get();
1615+
assertEquals(3, result);
1616+
List<Snapshot> snapshots = assertSnapshots(listener, 2, PROBE_ID1, PROBE_ID2);
1617+
// Snapshot 0
1618+
Snapshot snapshot0 = snapshots.get(0);
1619+
assertEquals(2, snapshot0.getCaptures().getReturn().getCaptureExpressions().size());
1620+
assertCaptureExpressions(
1621+
snapshot0.getCaptures().getReturn(),
1622+
"typed_fld_fld_msg",
1623+
String.class.getTypeName(),
1624+
"hello");
1625+
assertCaptureExpressions(
1626+
snapshot0.getCaptures().getReturn(), "nullTyped_fld", Object.class.getTypeName(), null);
1627+
assertNull(snapshot0.getCaptures().getReturn().getArguments());
1628+
assertNull(snapshot0.getCaptures().getReturn().getLocals());
1629+
// Snapshot 1
1630+
Snapshot snapshot1 = snapshots.get(1);
1631+
assertCaptureExpressions(
1632+
snapshot1.getCaptures().getReturn(), "var1", Integer.class.getTypeName(), "3");
1633+
assertCaptureExpressions(
1634+
snapshot1.getCaptures().getReturn(), "this_fld", Integer.class.getTypeName(), "11");
1635+
assertNull(snapshot1.getCaptures().getReturn().getArguments());
1636+
assertNull(snapshot1.getCaptures().getReturn().getLocals());
1637+
}
1638+
15251639
@Test
15261640
public void fields() throws IOException, URISyntaxException {
15271641
final String CLASS_NAME = "CapturedSnapshot06";

0 commit comments

Comments
 (0)