Skip to content

Commit edd5350

Browse files
evanchoolyevanchooly
committed
refactor the budget testing to better test the combos of line/session/capturing
1 parent c479c18 commit edd5350

1 file changed

Lines changed: 56 additions & 45 deletions

File tree

  • dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe

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

Lines changed: 56 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import datadog.trace.bootstrap.debugger.EvaluationError;
2323
import datadog.trace.bootstrap.debugger.MethodLocation;
2424
import datadog.trace.bootstrap.debugger.ProbeId;
25+
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
2526
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2627
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2728
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -30,6 +31,7 @@
3031
import datadog.trace.bootstrap.instrumentation.api.Tags;
3132
import datadog.trace.core.CoreTracer;
3233
import java.util.stream.Stream;
34+
import org.jetbrains.annotations.NotNull;
3335
import org.junit.jupiter.api.Assertions;
3436
import org.junit.jupiter.api.Test;
3537
import org.junit.jupiter.params.ParameterizedTest;
@@ -41,6 +43,7 @@ public class LogProbeTest {
4143
private static final String LANGUAGE = "java";
4244
private static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f8", 0);
4345
private static final String DEBUG_SESSION_ID = "TestSession";
46+
private static final int BUDGET_RUNS = 1100;
4447

4548
@Test
4649
public void testCapture() {
@@ -81,65 +84,68 @@ public void noDebugSession() {
8184

8285
@Test
8386
public void budgets() {
84-
BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
85-
DebuggerAgentHelper.injectSink(sink);
87+
try {
88+
ProbeRateLimiter.setGlobalSnapshotRate(-1);
89+
TracerAPI tracer =
90+
CoreTracer.builder()
91+
.idGenerationStrategy(IdGenerationStrategy.fromName("random"))
92+
.build();
93+
AgentTracer.registerIfAbsent(tracer);
94+
String sessionId = "12345";
8695

87-
TracerAPI tracer =
88-
CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
89-
AgentTracer.registerIfAbsent(tracer);
90-
int runs = 100;
91-
for (int i = 0; i < runs; i++) {
92-
runTrace(tracer, true);
93-
}
94-
assertEquals(runs * LogProbe.CAPTURING_PROBE_BUDGET, sink.captures);
96+
Result result = getResult(tracer, sessionId, true, null);
97+
assertEquals(BUDGET_RUNS * LogProbe.CAPTURING_PROBE_BUDGET, result.sink.captures);
9598

96-
sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
97-
DebuggerAgentHelper.injectSink(sink);
98-
runs = 1010;
99-
for (int i = 0; i < runs; i++) {
100-
runTrace(tracer, false);
101-
}
102-
assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
103-
}
99+
result = getResult(tracer, sessionId, false, null);
100+
assertEquals(BUDGET_RUNS * LogProbe.NON_CAPTURING_PROBE_BUDGET, result.sink.highRate);
104101

105-
@Test
106-
public void budgetsOnLineProbes() {
107-
BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
108-
DebuggerAgentHelper.injectSink(sink);
102+
// run without a session
103+
result = getResult(tracer, null, true, 100);
104+
assertEquals(result.count, result.sink.captures);
109105

110-
TracerAPI tracer =
111-
CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
112-
AgentTracer.registerIfAbsent(tracer);
113-
int runs = 100;
114-
for (int i = 0; i < runs; i++) {
115-
runTrace(tracer, true, 100);
106+
result = getResult(tracer, null, false, 100);
107+
assertEquals(result.count, result.sink.highRate);
108+
} finally {
109+
ProbeRateLimiter.resetGlobalRate();
116110
}
117-
assertEquals(runs * LogProbe.CAPTURING_PROBE_BUDGET, sink.captures);
111+
}
118112

119-
sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
113+
@NotNull
114+
private Result getResult(
115+
TracerAPI tracer, String sessionId, boolean captureSnapshot, Integer line) {
116+
BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
120117
DebuggerAgentHelper.injectSink(sink);
121-
runs = 1010;
122-
for (int i = 0; i < runs; i++) {
123-
runTrace(tracer, false, 100);
118+
int count = 0;
119+
for (int i = 0; i < BUDGET_RUNS; i++) {
120+
count += runTrace(tracer, captureSnapshot, line, sessionId);
124121
}
125-
assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
122+
return new Result(sink, count);
126123
}
127124

128-
private void runTrace(TracerAPI tracer, boolean captureSnapshot) {
129-
runTrace(tracer, captureSnapshot, null);
125+
private static class Result {
126+
final int count;
127+
final BudgetSink sink;
128+
129+
private Result(BudgetSink sink, int count) {
130+
this.sink = sink;
131+
this.count = count;
132+
}
130133
}
131134

132-
private void runTrace(TracerAPI tracer, boolean captureSnapshot, Integer line) {
135+
private int runTrace(TracerAPI tracer, boolean captureSnapshot, Integer line, String sessionId) {
133136
AgentSpan span = tracer.startSpan("budget testing", "test span");
134-
span.setTag(Tags.PROPAGATED_DEBUG, "12345:1");
137+
if (sessionId != null) {
138+
span.setTag(Tags.PROPAGATED_DEBUG, sessionId + ":1");
139+
}
135140
try (AgentScope scope = tracer.activateSpan(span, ScopeSource.MANUAL)) {
141+
Builder builder =
142+
createLog("Budget testing").probeId(ProbeId.newId()).captureSnapshot(captureSnapshot);
143+
if (sessionId != null) {
144+
builder.tags("session_id:" + sessionId);
145+
}
146+
LogProbe logProbe = builder.build();
147+
ProbeRateLimiter.setRate(logProbe.id, -1, captureSnapshot);
136148

137-
LogProbe logProbe =
138-
createLog("I'm in a debug session")
139-
.probeId(ProbeId.newId())
140-
.tags("session_id:12345")
141-
.captureSnapshot(captureSnapshot)
142-
.build();
143149
CapturedContext entryContext = capturedContext(span, logProbe);
144150
CapturedContext exitContext = capturedContext(span, logProbe);
145151
logProbe.evaluate(entryContext, new LogStatus(logProbe), MethodLocation.ENTRY);
@@ -158,7 +164,11 @@ private void runTrace(TracerAPI tracer, boolean captureSnapshot, Integer line) {
158164
logProbe.commit(entryContext, line);
159165
}
160166
}
161-
assertEquals(runs, span.getLocalRootSpan().getTag(format("_dd.ld.probe_id.%s", logProbe.id)));
167+
if (sessionId != null) {
168+
assertEquals(
169+
runs, span.getLocalRootSpan().getTag(format("_dd.ld.probe_id.%s", logProbe.id)));
170+
}
171+
return runs;
162172
}
163173
}
164174

@@ -343,6 +353,7 @@ private Builder createLog(String template) {
343353
private static class BudgetSink extends DebuggerSink {
344354

345355
public int captures;
356+
346357
public int highRate;
347358

348359
public BudgetSink(Config config, ProbeStatusSink probeStatusSink) {

0 commit comments

Comments
 (0)