add a tag tracking the overall invocation account

Author: evanchooly

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;
@@ -575,14 +576,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 +870,7 @@ public String toString() {
 
   private boolean inBudget() {
     AtomicInteger budgetLevel = getBudgetLevel();
-    return budgetLevel == null || budgetLevel.get() < PROBE_BUDGET;
+    return budgetLevel == null || budgetLevel.get() <= PROBE_BUDGET;
   }
 
   private AtomicInteger getBudgetLevel() {
@@ -881,6 +885,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/test/java/com/datadog/debugger/probe/LogProbeTest.java

@@ -110,9 +110,12 @@ private void runTrace(TracerAPI tracer) {
       logProbe.evaluate(entryContext, new LogStatus(logProbe), MethodLocation.ENTRY);
       logProbe.evaluate(exitContext, new LogStatus(logProbe), MethodLocation.EXIT);
 
-      for (int i = 0; i < 20; i++) {
+      for (int i = 0; i < LogProbe.PROBE_BUDGET * 2; i++) {
         logProbe.commit(entryContext, exitContext, emptyList());
       }
+      assertEquals(
+          LogProbe.PROBE_BUDGET * 2,
+          span.getLocalRootSpan().getTag(format("_dd.ld.probe_id.%s", logProbe.id)));
     }
   }