institute a 10 snapshot per probe per trace budget

Author: evanchooly

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

@@ -16,12 +16,14 @@
 import com.datadog.debugger.sink.DebuggerSink;
 import com.datadog.debugger.sink.Snapshot;
 import com.datadog.debugger.util.MoshiHelper;
+import com.datadog.debugger.util.WeakIdentityHashMap;
 import com.squareup.moshi.Json;
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.JsonReader;
 import com.squareup.moshi.JsonWriter;
 import com.squareup.moshi.Types;
 import datadog.trace.api.Config;
+import datadog.trace.api.DDTraceId;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.CorrelationAccess;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
@@ -49,6 +51,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -59,6 +62,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;
 
+  protected static final int PROBE_BUDGET = 10;
+
   /** Stores part of a templated message either a str or an expression */
   public static class Segment {
     private final String str;
@@ -278,6 +283,7 @@ public String toString() {
   private final Capture capture;
   private final Sampling sampling;
   private transient Consumer<Snapshot> snapshotProcessor;
+  private transient Map<DDTraceId, AtomicInteger> budget = new WeakIdentityHashMap<>();
 
   // no-arg constructor is required by Moshi to avoid creating instance with unsafe and by-passing
   // constructors, including field initializers.
@@ -568,7 +574,8 @@ public void commit(
       CapturedContext exitContext,
       List<CapturedContext.CapturedThrowable> caughtExceptions) {
     Snapshot snapshot = createSnapshot();
-    boolean shouldCommit = fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot);
+    boolean shouldCommit =
+        inBudget() && fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot);
     DebuggerSink sink = DebuggerAgent.getSink();
     if (shouldCommit) {
       commitSnapshot(snapshot, sink);
@@ -855,6 +862,15 @@ public String toString() {
     }
   }
 
+  private boolean inBudget() {
+    AgentSpan span = AgentTracer.activeSpan();
+    return span == null
+        || budget
+                .computeIfAbsent(span.getLocalRootSpan().getTraceId(), id -> new AtomicInteger())
+                .getAndIncrement()
+            < PROBE_BUDGET;
+  }
+
   @Generated
   @Override
   public boolean equals(Object o) {

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

@@ -53,6 +53,10 @@ 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

@@ -78,6 +78,44 @@ public void noDebugSession() {
         "With no debug sessions, snapshots should get filled.");
   }
 
+  @Test
+  public void budgets() {
+    DebuggerSink sink = new DebuggerSink(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);
+    }
+    assertEquals(runs * LogProbe.PROBE_BUDGET, sink.getSnapshotSink().getLowRateSnapshots().size());
+  }
+
+  private void runTrace(TracerAPI tracer) {
+    AgentSpan span = tracer.startSpan("budget testing", "test span");
+    try (AgentScope scope = tracer.activateSpan(span, ScopeSource.MANUAL)) {
+
+      LogProbe logProbe =
+          createLog("I'm in a debug session")
+              .probeId(ProbeId.newId())
+              .captureSnapshot(true)
+              .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++) {
+        logProbe.commit(entryContext, exitContext, emptyList());
+      }
+      tracer.startSpan("budget testing", "test span");
+      tracer.activateNext(span);
+    }
+  }
+
   private boolean fillSnapshot(DebugSessionStatus status) {
     DebuggerAgentHelper.injectSink(new DebuggerSink(getConfig(), mock(ProbeStatusSink.class)));
     TracerAPI tracer =