Use two budgets depending on type (#8283)

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();

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<CapturedContext.CapturedThrowable> 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());
+      }
     }
   }
 

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<Snapshot> getLowRateSnapshots() {
-    return lowRateSnapshots;
-  }
-
   public void start() {
     if (started.compareAndSet(false, true)) {
       highRateScheduled =

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();
+    }
+  }
 }