diff --git a/dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java b/dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java index 9d1c793ea8e..5acd251b78a 100644 --- a/dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java +++ b/dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java @@ -39,6 +39,12 @@ public String tag() { public String tag() { return "cause:debug session disabled"; } + }, + BUDGET { + @Override + public String tag() { + return "cause:budget_exceeded"; + } }; public abstract String tag(); diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java index f07f17cc558..95ba3426135 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java +++ b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java @@ -1,6 +1,7 @@ package com.datadog.debugger.probe; import static com.datadog.debugger.probe.LogProbe.Capture.toLimits; +import static java.lang.String.format; import com.datadog.debugger.agent.DebuggerAgent; import com.datadog.debugger.agent.Generated; @@ -62,7 +63,8 @@ public class LogProbe extends ProbeDefinition implements Sampled { private static final Limits LIMITS = new Limits(1, 3, 8192, 5); private static final int LOG_MSG_LIMIT = 8192; - public static final int PROBE_BUDGET = 10; + public static final int CAPTURING_PROBE_BUDGET = 10; + public static final int NON_CAPTURING_PROBE_BUDGET = 1000; /** Stores part of a templated message either a str or an expression */ public static class Segment { @@ -575,14 +577,17 @@ public void commit( CapturedContext exitContext, List caughtExceptions) { Snapshot snapshot = createSnapshot(); - boolean shouldCommit = - inBudget() && fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot); + boolean shouldCommit = fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot); DebuggerSink sink = DebuggerAgent.getSink(); if (shouldCommit) { - commitSnapshot(snapshot, sink); incrementBudget(); - if (snapshotProcessor != null) { - snapshotProcessor.accept(snapshot); + if (inBudget()) { + commitSnapshot(snapshot, sink); + if (snapshotProcessor != null) { + snapshotProcessor.accept(snapshot); + } + } else { + sink.skipSnapshot(id, DebuggerContext.SkipCause.BUDGET); } } else { sink.skipSnapshot(id, DebuggerContext.SkipCause.CONDITION); @@ -866,7 +871,9 @@ public String toString() { private boolean inBudget() { AtomicInteger budgetLevel = getBudgetLevel(); - return budgetLevel == null || budgetLevel.get() < PROBE_BUDGET; + return budgetLevel == null + || budgetLevel.get() + <= (captureSnapshot ? CAPTURING_PROBE_BUDGET : NON_CAPTURING_PROBE_BUDGET); } private AtomicInteger getBudgetLevel() { @@ -881,6 +888,11 @@ private void incrementBudget() { AtomicInteger budgetLevel = getBudgetLevel(); if (budgetLevel != null) { budgetLevel.incrementAndGet(); + TracerAPI tracer = AgentTracer.get(); + AgentSpan span = tracer != null ? tracer.activeSpan() : null; + if (span != null) { + span.getLocalRootSpan().setTag(format("_dd.ld.probe_id.%s", id), budgetLevel.get()); + } } } diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SnapshotSink.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SnapshotSink.java index 168ef43a19e..8b47882c8fa 100644 --- a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SnapshotSink.java +++ b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SnapshotSink.java @@ -53,10 +53,6 @@ public SnapshotSink(Config config, String tags, BatchUploader snapshotUploader) this.snapshotUploader = snapshotUploader; } - public BlockingQueue getLowRateSnapshots() { - return lowRateSnapshots; - } - public void start() { if (started.compareAndSet(false, true)) { highRateScheduled = diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java index 19019ffa3e6..5c048a47fd8 100644 --- a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java +++ b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/probe/LogProbeTest.java @@ -16,6 +16,7 @@ import com.datadog.debugger.sink.DebuggerSink; import com.datadog.debugger.sink.ProbeStatusSink; import com.datadog.debugger.sink.Snapshot; +import datadog.trace.api.Config; import datadog.trace.api.IdGenerationStrategy; import datadog.trace.bootstrap.debugger.CapturedContext; import datadog.trace.bootstrap.debugger.EvaluationError; @@ -80,20 +81,28 @@ public void noDebugSession() { @Test public void budgets() { - DebuggerSink sink = new DebuggerSink(getConfig(), mock(ProbeStatusSink.class)); + BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class)); DebuggerAgentHelper.injectSink(sink); - assertEquals(0, sink.getSnapshotSink().getLowRateSnapshots().size()); + TracerAPI tracer = CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build(); AgentTracer.registerIfAbsent(tracer); int runs = 100; for (int i = 0; i < runs; i++) { - runTrace(tracer); + runTrace(tracer, true); + } + assertEquals(runs * LogProbe.CAPTURING_PROBE_BUDGET, sink.captures); + + sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class)); + DebuggerAgentHelper.injectSink(sink); + runs = 1010; + for (int i = 0; i < runs; i++) { + runTrace(tracer, false); } - assertEquals(runs * LogProbe.PROBE_BUDGET, sink.getSnapshotSink().getLowRateSnapshots().size()); + assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate); } - private void runTrace(TracerAPI tracer) { + private void runTrace(TracerAPI tracer, boolean captureSnapshot) { AgentSpan span = tracer.startSpan("budget testing", "test span"); span.setTag(Tags.PROPAGATED_DEBUG, "12345:1"); try (AgentScope scope = tracer.activateSpan(span, ScopeSource.MANUAL)) { @@ -102,17 +111,23 @@ private void runTrace(TracerAPI tracer) { createLog("I'm in a debug session") .probeId(ProbeId.newId()) .tags("session_id:12345") - .captureSnapshot(true) + .captureSnapshot(captureSnapshot) .build(); - CapturedContext entryContext = capturedContext(span, logProbe); CapturedContext exitContext = capturedContext(span, logProbe); logProbe.evaluate(entryContext, new LogStatus(logProbe), MethodLocation.ENTRY); logProbe.evaluate(exitContext, new LogStatus(logProbe), MethodLocation.EXIT); - for (int i = 0; i < 20; i++) { + int budget = + logProbe.isCaptureSnapshot() + ? LogProbe.CAPTURING_PROBE_BUDGET + : LogProbe.NON_CAPTURING_PROBE_BUDGET; + int runs = budget + 20; + + for (int i = 0; i < runs; i++) { logProbe.commit(entryContext, exitContext, emptyList()); } + assertEquals(runs, span.getLocalRootSpan().getTag(format("_dd.ld.probe_id.%s", logProbe.id))); } } @@ -293,4 +308,34 @@ private Builder createLog(String template) { .probeId(PROBE_ID) .template(template, parseTemplate(template)); } + + private static class BudgetSink extends DebuggerSink { + + public int captures; + public int highRate; + + public BudgetSink(Config config, ProbeStatusSink probeStatusSink) { + super(config, probeStatusSink); + } + + @Override + public void addHighRateSnapshot(Snapshot snapshot) { + highRate++; + } + + @Override + public void addSnapshot(Snapshot snapshot) { + captures++; + } + + @Override + public void start() { + super.start(); + } + + @Override + public void stop() { + super.stop(); + } + } }