make sure budgets apply to line probes, too.

Author: evanchooly

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

@@ -701,15 +701,19 @@ public void commit(CapturedContext lineContext, int line) {
       shouldCommit = true;
     }
     if (shouldCommit) {
-      if (isCaptureSnapshot()) {
-        // freeze context just before commit because line probes have only one context
-        Duration timeout =
-            Duration.of(Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
-        lineContext.freeze(new TimeoutChecker(timeout));
-        snapshot.addLine(lineContext, line);
+      incrementBudget();
+      if (inBudget()) {
+        if (isCaptureSnapshot()) {
+          // freeze context just before commit because line probes have only one context
+          Duration timeout =
+              Duration.of(
+                  Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
+          lineContext.freeze(new TimeoutChecker(timeout));
+          snapshot.addLine(lineContext, line);
+        }
+        commitSnapshot(snapshot, sink);
+        return;
       }
-      commitSnapshot(snapshot, sink);
-      return;
     }
     sink.skipSnapshot(id, DebuggerContext.SkipCause.CONDITION);
   }

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

@@ -102,7 +102,34 @@ public void budgets() {
     assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
   }
 
+  @Test
+  public void budgetsOnLineProbes() {
+    BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
+    DebuggerAgentHelper.injectSink(sink);
+
+    TracerAPI tracer =
+        CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
+    AgentTracer.registerIfAbsent(tracer);
+    int runs = 100;
+    for (int i = 0; i < runs; i++) {
+      runTrace(tracer, true, 100);
+    }
+    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, 100);
+    }
+    assertEquals(runs * LogProbe.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
+  }
+
   private void runTrace(TracerAPI tracer, boolean captureSnapshot) {
+    runTrace(tracer, captureSnapshot, null);
+  }
+
+  private void runTrace(TracerAPI tracer, boolean captureSnapshot, Integer line) {
     AgentSpan span = tracer.startSpan("budget testing", "test span");
     span.setTag(Tags.PROPAGATED_DEBUG, "12345:1");
     try (AgentScope scope = tracer.activateSpan(span, ScopeSource.MANUAL)) {
@@ -125,7 +152,11 @@ private void runTrace(TracerAPI tracer, boolean captureSnapshot) {
       int runs = budget + 20;
 
       for (int i = 0; i < runs; i++) {
-        logProbe.commit(entryContext, exitContext, emptyList());
+        if (line == null) {
+          logProbe.commit(entryContext, exitContext, emptyList());
+        } else {
+          logProbe.commit(entryContext, line);
+        }
       }
       assertEquals(runs, span.getLocalRootSpan().getTag(format("_dd.ld.probe_id.%s", logProbe.id)));
     }
@@ -172,6 +203,29 @@ private static CapturedContext capturedContext(AgentSpan span, ProbeDefinition p
     return context;
   }
 
+  @Test
+  public void testDebugLineProbes() {
+    BudgetSink sink = new BudgetSink(getConfig(), mock(ProbeStatusSink.class));
+    DebuggerAgentHelper.injectSink(sink);
+
+    TracerAPI tracer =
+        CoreTracer.builder().idGenerationStrategy(IdGenerationStrategy.fromName("random")).build();
+    AgentTracer.registerIfAbsent(tracer);
+    int runs = 100;
+    for (int i = 0; i < runs; i++) {
+      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.NON_CAPTURING_PROBE_BUDGET, sink.highRate);
+  }
+
   @Test
   public void log() {
     LogProbe logProbe = createLog(null).build();