Skip to content

Commit 71309b3

Browse files
evanchoolyevanchooly
authored
Institute a 10 snapshot per probe per trace budget (#8277)
1 parent 93b44ea commit 71309b3

File tree

3 files changed

+91
-21
lines changed

3 files changed

+91
-21
lines changed

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616
import com.datadog.debugger.sink.DebuggerSink;
1717
import com.datadog.debugger.sink.Snapshot;
1818
import com.datadog.debugger.util.MoshiHelper;
19+
import com.datadog.debugger.util.WeakIdentityHashMap;
1920
import com.squareup.moshi.Json;
2021
import com.squareup.moshi.JsonAdapter;
2122
import com.squareup.moshi.JsonReader;
2223
import com.squareup.moshi.JsonWriter;
2324
import com.squareup.moshi.Types;
2425
import datadog.trace.api.Config;
26+
import datadog.trace.api.DDTraceId;
2527
import datadog.trace.bootstrap.debugger.CapturedContext;
2628
import datadog.trace.bootstrap.debugger.CorrelationAccess;
2729
import datadog.trace.bootstrap.debugger.DebuggerContext;
@@ -49,6 +51,7 @@
4951
import java.util.List;
5052
import java.util.Map;
5153
import java.util.Objects;
54+
import java.util.concurrent.atomic.AtomicInteger;
5255
import java.util.function.Consumer;
5356
import org.slf4j.Logger;
5457
import org.slf4j.LoggerFactory;
@@ -59,6 +62,8 @@ public class LogProbe extends ProbeDefinition implements Sampled {
5962
private static final Limits LIMITS = new Limits(1, 3, 8192, 5);
6063
private static final int LOG_MSG_LIMIT = 8192;
6164

65+
public static final int PROBE_BUDGET = 10;
66+
6267
/** Stores part of a templated message either a str or an expression */
6368
public static class Segment {
6469
private final String str;
@@ -278,6 +283,8 @@ public String toString() {
278283
private final Capture capture;
279284
private final Sampling sampling;
280285
private transient Consumer<Snapshot> snapshotProcessor;
286+
protected transient Map<DDTraceId, AtomicInteger> budget =
287+
Collections.synchronizedMap(new WeakIdentityHashMap<>());
281288

282289
// no-arg constructor is required by Moshi to avoid creating instance with unsafe and by-passing
283290
// constructors, including field initializers.
@@ -568,10 +575,12 @@ public void commit(
568575
CapturedContext exitContext,
569576
List<CapturedContext.CapturedThrowable> caughtExceptions) {
570577
Snapshot snapshot = createSnapshot();
571-
boolean shouldCommit = fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot);
578+
boolean shouldCommit =
579+
inBudget() && fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot);
572580
DebuggerSink sink = DebuggerAgent.getSink();
573581
if (shouldCommit) {
574582
commitSnapshot(snapshot, sink);
583+
incrementBudget();
575584
if (snapshotProcessor != null) {
576585
snapshotProcessor.accept(snapshot);
577586
}
@@ -855,6 +864,26 @@ public String toString() {
855864
}
856865
}
857866

867+
private boolean inBudget() {
868+
AtomicInteger budgetLevel = getBudgetLevel();
869+
return budgetLevel == null || budgetLevel.get() < PROBE_BUDGET;
870+
}
871+
872+
private AtomicInteger getBudgetLevel() {
873+
TracerAPI tracer = AgentTracer.get();
874+
AgentSpan span = tracer != null ? tracer.activeSpan() : null;
875+
return getDebugSessionId() == null || span == null
876+
? null
877+
: budget.computeIfAbsent(span.getLocalRootSpan().getTraceId(), id -> new AtomicInteger());
878+
}
879+
880+
private void incrementBudget() {
881+
AtomicInteger budgetLevel = getBudgetLevel();
882+
if (budgetLevel != null) {
883+
budgetLevel.incrementAndGet();
884+
}
885+
}
886+
858887
@Generated
859888
@Override
860889
public boolean equals(Object o) {

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ public SnapshotSink(Config config, String tags, BatchUploader snapshotUploader)
5353
this.snapshotUploader = snapshotUploader;
5454
}
5555

56+
public BlockingQueue<Snapshot> getLowRateSnapshots() {
57+
return lowRateSnapshots;
58+
}
59+
5660
public void start() {
5761
if (started.compareAndSet(false, true)) {
5862
highRateScheduled =

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

Lines changed: 57 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -43,18 +43,18 @@ public class LogProbeTest {
4343

4444
@Test
4545
public void testCapture() {
46-
LogProbe.Builder builder = createLog(null);
46+
Builder builder = createLog(null);
4747
LogProbe snapshotProbe = builder.capture(1, 420, 255, 20).build();
48-
Assertions.assertEquals(1, snapshotProbe.getCapture().getMaxReferenceDepth());
49-
Assertions.assertEquals(420, snapshotProbe.getCapture().getMaxCollectionSize());
50-
Assertions.assertEquals(255, snapshotProbe.getCapture().getMaxLength());
48+
assertEquals(1, snapshotProbe.getCapture().getMaxReferenceDepth());
49+
assertEquals(420, snapshotProbe.getCapture().getMaxCollectionSize());
50+
assertEquals(255, snapshotProbe.getCapture().getMaxLength());
5151
}
5252

5353
@Test
5454
public void testSampling() {
55-
LogProbe.Builder builder = createLog(null);
55+
Builder builder = createLog(null);
5656
LogProbe snapshotProbe = builder.sampling(0.25).build();
57-
Assertions.assertEquals(0.25, snapshotProbe.getSampling().getEventsPerSecond(), 0.01);
57+
assertEquals(0.25, snapshotProbe.getSampling().getEventsPerSecond(), 0.01);
5858
}
5959

6060
@Test
@@ -78,6 +78,44 @@ public void noDebugSession() {
7878
"With no debug sessions, snapshots should get filled.");
7979
}
8080

81+
@Test
82+
public void budgets() {
83+
DebuggerSink sink = new DebuggerSink(getConfig(), mock(ProbeStatusSink.class));
84+
DebuggerAgentHelper.injectSink(sink);
85+
assertEquals(0, sink.getSnapshotSink().getLowRateSnapshots().size());
86+
TracerAPI tracer =
87+
CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
88+
AgentTracer.registerIfAbsent(tracer);
89+
int runs = 100;
90+
for (int i = 0; i < runs; i++) {
91+
runTrace(tracer);
92+
}
93+
assertEquals(runs * LogProbe.PROBE_BUDGET, sink.getSnapshotSink().getLowRateSnapshots().size());
94+
}
95+
96+
private void runTrace(TracerAPI tracer) {
97+
AgentSpan span = tracer.startSpan("budget testing", "test span");
98+
span.setTag(Tags.PROPAGATED_DEBUG, "12345:1");
99+
try (AgentScope scope = tracer.activateSpan(span, ScopeSource.MANUAL)) {
100+
101+
LogProbe logProbe =
102+
createLog("I'm in a debug session")
103+
.probeId(ProbeId.newId())
104+
.tags("session_id:12345")
105+
.captureSnapshot(true)
106+
.build();
107+
108+
CapturedContext entryContext = capturedContext(span, logProbe);
109+
CapturedContext exitContext = capturedContext(span, logProbe);
110+
logProbe.evaluate(entryContext, new LogStatus(logProbe), MethodLocation.ENTRY);
111+
logProbe.evaluate(exitContext, new LogStatus(logProbe), MethodLocation.EXIT);
112+
113+
for (int i = 0; i < 20; i++) {
114+
logProbe.commit(entryContext, exitContext, emptyList());
115+
}
116+
}
117+
}
118+
81119
private boolean fillSnapshot(DebugSessionStatus status) {
82120
DebuggerAgentHelper.injectSink(new DebuggerSink(getConfig(), mock(ProbeStatusSink.class)));
83121
TracerAPI tracer =
@@ -152,12 +190,11 @@ public void fillSnapshot_shouldSend(String methodLocation) {
152190
LogProbe logProbe = createLog(null).evaluateAt(MethodLocation.valueOf(methodLocation)).build();
153191
CapturedContext entryContext = new CapturedContext();
154192
CapturedContext exitContext = new CapturedContext();
155-
LogProbe.LogStatus logEntryStatus =
156-
prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
193+
LogStatus logEntryStatus = prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
157194
logEntryStatus.setSampled(true); // force sampled to avoid rate limiting executing tests!
158-
LogProbe.LogStatus logExitStatus = prepareContext(exitContext, logProbe, MethodLocation.EXIT);
195+
LogStatus logExitStatus = prepareContext(exitContext, logProbe, MethodLocation.EXIT);
159196
logExitStatus.setSampled(true); // force sampled to avoid rate limiting executing tests!
160-
Snapshot snapshot = new Snapshot(Thread.currentThread(), logProbe, 10);
197+
Snapshot snapshot = new Snapshot(currentThread(), logProbe, 10);
161198
assertTrue(logProbe.fillSnapshot(entryContext, exitContext, null, snapshot));
162199
}
163200

@@ -172,16 +209,16 @@ public void fillSnapshot(
172209
LogProbe logProbe = createLog(null).evaluateAt(MethodLocation.EXIT).build();
173210
CapturedContext entryContext = new CapturedContext();
174211
CapturedContext exitContext = new CapturedContext();
175-
LogProbe.LogStatus entryStatus = prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
212+
LogStatus entryStatus = prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
176213
fillStatus(entryStatus, sampled, condition, conditionErrors, logTemplateErrors);
177-
LogProbe.LogStatus exitStatus = prepareContext(exitContext, logProbe, MethodLocation.EXIT);
214+
LogStatus exitStatus = prepareContext(exitContext, logProbe, MethodLocation.EXIT);
178215
fillStatus(exitStatus, sampled, condition, conditionErrors, logTemplateErrors);
179-
Snapshot snapshot = new Snapshot(Thread.currentThread(), logProbe, 10);
216+
Snapshot snapshot = new Snapshot(currentThread(), logProbe, 10);
180217
assertEquals(shouldCommit, logProbe.fillSnapshot(entryContext, exitContext, null, snapshot));
181218
}
182219

183220
private void fillStatus(
184-
LogProbe.LogStatus entryStatus,
221+
LogStatus entryStatus,
185222
boolean sampled,
186223
boolean condition,
187224
boolean conditionErrors,
@@ -193,10 +230,10 @@ private void fillStatus(
193230
entryStatus.setLogTemplateErrors(logTemplateErrors);
194231
}
195232

196-
private LogProbe.LogStatus prepareContext(
233+
private LogStatus prepareContext(
197234
CapturedContext context, LogProbe logProbe, MethodLocation methodLocation) {
198235
context.evaluate(PROBE_ID.getEncodedId(), logProbe, "", 0, methodLocation);
199-
return (LogProbe.LogStatus) context.getStatus(PROBE_ID.getEncodedId());
236+
return (LogStatus) context.getStatus(PROBE_ID.getEncodedId());
200237
}
201238

202239
private static Stream<Arguments> statusValues() {
@@ -222,15 +259,15 @@ public void fillSnapshot_shouldSend_exit() {
222259
prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
223260
CapturedContext exitContext = new CapturedContext();
224261
prepareContext(exitContext, logProbe, MethodLocation.EXIT);
225-
Snapshot snapshot = new Snapshot(Thread.currentThread(), logProbe, 10);
262+
Snapshot snapshot = new Snapshot(currentThread(), logProbe, 10);
226263
assertTrue(logProbe.fillSnapshot(entryContext, exitContext, null, snapshot));
227264
}
228265

229266
@Test
230267
public void fillSnapshot_shouldSend_evalErrors() {
231268
LogProbe logProbe = createLog(null).evaluateAt(MethodLocation.EXIT).build();
232269
CapturedContext entryContext = new CapturedContext();
233-
LogProbe.LogStatus logStatus = prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
270+
LogStatus logStatus = prepareContext(entryContext, logProbe, MethodLocation.ENTRY);
234271
logStatus.addError(new EvaluationError("expr", "msg1"));
235272
logStatus.setLogTemplateErrors(true);
236273
entryContext.addThrowable(new RuntimeException("errorEntry"));
@@ -239,7 +276,7 @@ public void fillSnapshot_shouldSend_evalErrors() {
239276
logStatus.addError(new EvaluationError("expr", "msg2"));
240277
logStatus.setLogTemplateErrors(true);
241278
exitContext.addThrowable(new RuntimeException("errorExit"));
242-
Snapshot snapshot = new Snapshot(Thread.currentThread(), logProbe, 10);
279+
Snapshot snapshot = new Snapshot(currentThread(), logProbe, 10);
243280
assertTrue(logProbe.fillSnapshot(entryContext, exitContext, null, snapshot));
244281
assertEquals(2, snapshot.getEvaluationErrors().size());
245282
assertEquals("msg1", snapshot.getEvaluationErrors().get(0).getMessage());
@@ -250,7 +287,7 @@ public void fillSnapshot_shouldSend_evalErrors() {
250287
"errorExit", snapshot.getCaptures().getReturn().getCapturedThrowable().getMessage());
251288
}
252289

253-
private LogProbe.Builder createLog(String template) {
290+
private Builder createLog(String template) {
254291
return LogProbe.builder()
255292
.language(LANGUAGE)
256293
.probeId(PROBE_ID)

0 commit comments

Comments
 (0)