Add coordinated sampling for snapshot probes (#12452)

Add coordinated sampling for snapshot probes

Adds coordinated sampling for full debugger snapshots using the active
Datadog context, using Context API.
State of Coordinated Sampling is stored into a Context attached to
the root local span.
The first probe’s sampling decision controls related probes, avoiding
fragmented snapshot sets.
Ensures each probe emits at most once per coordinated context.
Log-only probes remain independent.
Encapsulates ProbeDefinition.probeId, updating callers to use
getProbeId().

address codex comments

fix exception probe sampling

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: jpbempel

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

@@ -126,6 +126,6 @@ public int hashCode() {
   public String toString() {
     return String.format(
         "CodeOriginProbe{probeId=%s, entrySpanProbe=%s, signature=%s, where=%s, location=%s}",
-        probeId, entrySpanProbe, signature, where, location);
+        getProbeId(), entrySpanProbe, signature, where, location);
   }
 }

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

@@ -68,6 +68,13 @@ public boolean isLineProbe() {
     return false;
   }
 
+  @Override
+  protected boolean useCoordinatedSampling() {
+    // exception probes have their own independent exception sampling flow and must not share a
+    // trace-wide sampling decision with ordinary snapshot probes on the same local root span.
+    return false;
+  }
+
   @Override
   public CapturedContext.Status createStatus() {
     return new ExceptionProbeStatus(this);

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

@@ -1,6 +1,8 @@
 package com.datadog.debugger.probe;
 
 import static com.datadog.debugger.probe.LogProbe.Capture.toLimits;
+import static com.datadog.debugger.probe.LogProbe.CoordinatedSamplingState.Status.DROP;
+import static com.datadog.debugger.probe.LogProbe.CoordinatedSamplingState.Status.EMIT;
 import static datadog.trace.api.debugger.DebuggerMetricCollector.SkippedReason.RATE_LIMIT;
 import static java.lang.String.format;
 
@@ -24,6 +26,8 @@
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.JsonReader;
 import com.squareup.moshi.JsonWriter;
+import datadog.context.Context;
+import datadog.context.ContextKey;
 import datadog.trace.api.Config;
 import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.DDTraceId;
@@ -52,6 +56,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
@@ -62,7 +68,8 @@ public class LogProbe extends ProbeDefinition implements Sampled, CapturedContex
   private static final Logger LOGGER = LoggerFactory.getLogger(LogProbe.class);
   private static final Limits LIMITS = new Limits(1, 3, 8192, 5);
   private static final int LOG_MSG_LIMIT = 8192;
-
+  private static final ContextKey<CoordinatedSamplingState> SAMPLING_KEY =
+      ContextKey.named("debugger_sampling");
   public static final int CAPTURING_PROBE_BUDGET = 10;
   public static final int NON_CAPTURING_PROBE_BUDGET = 1000;
 
@@ -511,7 +518,7 @@ public InstrumentationResult.Status instrument(
   public boolean isReadyToCapture() {
     if (!hasCondition()) {
       // we are sampling here to avoid creating CapturedContext when the sampling result is negative
-      boolean sampled = ProbeRateLimiter.tryProbe(sampler, isFullSnapshot());
+      boolean sampled = trySample(sampler);
       if (!sampled) {
         DebuggerAgent.getSink().skipSnapshot(id, RATE_LIMIT);
       }
@@ -581,15 +588,29 @@ private void sample(LogStatus logStatus, MethodLocation methodLocation) {
     if (!MethodLocation.isSame(methodLocation, evaluateAt)) {
       return;
     }
-    // if condition has error and no capture Snapshot, the error is reported using errorSampler
-    // at 1/s rate instead of the log template one
-    Sampler localSampler =
-        logStatus.hasConditionErrors && !isFullSnapshot() ? errorSampler : sampler;
-    boolean sampled =
-        !logStatus.getDebugSessionStatus().isDisabled()
-            && ProbeRateLimiter.tryProbe(localSampler, isFullSnapshot());
+    DebugSessionStatus debugSessionStatus = logStatus.getDebugSessionStatus();
+    if (debugSessionStatus.isDisabled()) {
+      return;
+    }
+    boolean sampled;
+    if (debugSessionStatus.isActive()) {
+      // LogStatus.shouldSend() emits unconditionally when the debug session is ACTIVE, so this
+      // probe must not record a coordinated decision that shouldSend() would ignore anyway.
+      // Instead, make sure the shared trace-level state lets the other full-snapshot probes on
+      // this trace emit too, rather than leaving/recording a DROP that would suppress them.
+      sampled = true;
+      if (isFullSnapshot() && useCoordinatedSampling()) {
+        forceCoordinatedEmit();
+      }
+    } else {
+      // if condition has error and no capture Snapshot, the error is reported using errorSampler
+      // at 1/s rate instead of the log template one
+      Sampler localSampler =
+          logStatus.hasConditionErrors && !isFullSnapshot() ? errorSampler : sampler;
+      sampled = trySample(localSampler);
+    }
     logStatus.setSampled(sampled);
-    if (!sampled && !logStatus.getDebugSessionStatus().isDisabled()) {
+    if (!sampled) {
       DebuggerAgent.getSink().skipSnapshot(id, RATE_LIMIT);
     }
   }
@@ -642,6 +663,104 @@ public void commit(
     }
   }
 
+  /**
+   * Holds the once-per-trace coordinated sampling decision shared by all full-snapshot probes on a
+   * given local root span, so that either all of them emit or none of them do.
+   */
+  static class CoordinatedSamplingState {
+    /** Outcome of the first probe's sampling decision for the trace, cached on the root span. */
+    enum Status {
+      /** The trace was not sampled; every probe sharing this state must not emit. */
+      DROP,
+      /** The trace was sampled; probes sharing this state may emit, once each. */
+      EMIT
+    }
+
+    private final Set<String> emittedProbeIds;
+    private final Status status;
+
+    CoordinatedSamplingState(Status status) {
+      this.status = status;
+      if (status == EMIT) {
+        emittedProbeIds = ConcurrentHashMap.newKeySet();
+      } else {
+        emittedProbeIds = Collections.emptySet();
+      }
+    }
+
+    boolean tryEmit(String probeEncodedId) {
+      return status == EMIT && emittedProbeIds.add(probeEncodedId);
+    }
+  }
+
+  /**
+   * Whether this probe shares the trace-wide {@link CoordinatedSamplingState} with other
+   * full-snapshot probes on the same local root span, so they all emit together or not at all.
+   * Probes with their own independent sampling flow (e.g. {@link ExceptionProbe}) must not
+   * participate, or their unrelated decision would leak into (and be leaked into by) ordinary
+   * snapshot probes sharing the same trace.
+   */
+  protected boolean useCoordinatedSampling() {
+    return true;
+  }
+
+  private boolean trySample(Sampler sampler) {
+    if (!isFullSnapshot()) {
+      return ProbeRateLimiter.tryProbe(sampler, false);
+    }
+    if (!useCoordinatedSampling()) {
+      return ProbeRateLimiter.tryProbe(sampler, true);
+    }
+
+    AgentSpan localRootSpan = getActiveLocalRootSpan();
+    if (localRootSpan == null) {
+      return ProbeRateLimiter.tryProbe(sampler, true);
+    }
+
+    CoordinatedSamplingState state = Context.from(localRootSpan).get(SAMPLING_KEY);
+    if (state == null) {
+      synchronized (localRootSpan) {
+        Context traceContext = Context.from(localRootSpan);
+        state = traceContext.get(SAMPLING_KEY);
+        if (state == null) {
+          boolean sampled = ProbeRateLimiter.tryProbe(sampler, true);
+          state = new CoordinatedSamplingState(sampled ? EMIT : DROP);
+          traceContext.with(SAMPLING_KEY, state).attachTo(localRootSpan);
+        }
+      }
+    }
+    return state.tryEmit(getProbeId().getEncodedId());
+  }
+
+  /**
+   * Ensures the trace-level coordinated sampling state allows emission, upgrading a previously
+   * cached DROP (recorded by an ordinary probe before this debug session became active) so that
+   * every full-snapshot probe on this trace gets a chance to emit.
+   */
+  private void forceCoordinatedEmit() {
+    AgentSpan localRootSpan = getActiveLocalRootSpan();
+    if (localRootSpan == null) {
+      return;
+    }
+    synchronized (localRootSpan) {
+      Context traceContext = Context.from(localRootSpan);
+      CoordinatedSamplingState state = traceContext.get(SAMPLING_KEY);
+      if (state == null || state.status != EMIT) {
+        traceContext.with(SAMPLING_KEY, new CoordinatedSamplingState(EMIT)).attachTo(localRootSpan);
+      }
+    }
+  }
+
+  private AgentSpan getActiveLocalRootSpan() {
+    TracerAPI tracer = AgentTracer.get();
+    AgentSpan activeSpan = tracer != null ? tracer.activeSpan() : null;
+    if (activeSpan == null || activeSpan == AgentTracer.noopSpan()) {
+      return null;
+    }
+    AgentSpan localRootSpan = activeSpan.getLocalRootSpan();
+    return localRootSpan != AgentTracer.noopSpan() ? localRootSpan : null;
+  }
+
   protected Snapshot createSnapshot() {
     int maxDepth = capture != null ? capture.maxReferenceDepth : -1;
     return new Snapshot(Thread.currentThread(), this, maxDepth);
@@ -652,8 +771,9 @@ protected boolean fillSnapshot(
       CapturedContext exitContext,
       List<CapturedContext.CapturedThrowable> caughtExceptions,
       Snapshot snapshot) {
-    LogStatus entryStatus = convertStatus(entryContext.getStatus(probeId.getEncodedId()));
-    LogStatus exitStatus = convertStatus(exitContext.getStatus(probeId.getEncodedId()));
+    String probeEncodedId = getProbeId().getEncodedId();
+    LogStatus entryStatus = convertStatus(entryContext.getStatus(probeEncodedId));
+    LogStatus exitStatus = convertStatus(exitContext.getStatus(probeEncodedId));
     String message = null;
     switch (evaluateAt) {
       case ENTRY:
@@ -829,7 +949,7 @@ protected void commitSnapshot(Snapshot snapshot, DebuggerSink sink) {
 
   @Override
   public void commit(CapturedContext lineContext, int line) {
-    LogStatus status = (LogStatus) lineContext.getStatus(probeId.getEncodedId());
+    LogStatus status = (LogStatus) lineContext.getStatus(getProbeId().getEncodedId());
     if (status == null) {
       return;
     }

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

@@ -28,7 +28,7 @@ public abstract class ProbeDefinition implements ProbeImplementation {
   protected final String language;
   protected final String id;
   protected final int version;
-  protected transient ProbeId probeId;
+  private transient ProbeId probeId;
   protected final Tag[] tags;
   protected final Map<String, String> tagMap = new HashMap<>();
   protected final Where where;

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

@@ -273,10 +273,11 @@ public void commit(
       CapturedContext entryContext,
       CapturedContext exitContext,
       List<CapturedContext.CapturedThrowable> caughtExceptions) {
+    String probeEncodedId = getProbeId().getEncodedId();
     CapturedContext.Status status =
         evaluateAt == MethodLocation.EXIT
-            ? exitContext.getStatus(probeId.getEncodedId())
-            : entryContext.getStatus(probeId.getEncodedId());
+            ? exitContext.getStatus(probeEncodedId)
+            : entryContext.getStatus(probeEncodedId);
     if (status == null) {
       return;
     }
@@ -287,7 +288,7 @@ public void commit(
 
   @Override
   public void commit(CapturedContext lineContext, int line) {
-    CapturedContext.Status status = lineContext.getStatus(probeId.getEncodedId());
+    CapturedContext.Status status = lineContext.getStatus(getProbeId().getEncodedId());
     if (status == null) {
       return;
     }
@@ -319,7 +320,7 @@ private void decorateTags(SpanDecorationStatus status) {
     }
     if (!tagsToDecorate.isEmpty()) {
       // only send EMITTING status if we set at least one tag
-      DebuggerAgent.getSink().getProbeStatusSink().addEmitting(probeId);
+      DebuggerAgent.getSink().getProbeStatusSink().addEmitting(getProbeId());
     }
   }
 

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

@@ -140,7 +140,7 @@ private boolean evaluateCondition(CapturedContext capture) {
       Duration timeout = Duration.ofMillis(Config.get().getDynamicInstrumentationEvalTimeout());
       return probeCondition.execute(capture, TimeoutChecker.create(Config.get(), timeout));
     } catch (Exception ex) {
-      DebuggerAgent.getSink().getProbeStatusSink().addError(probeId, ex);
+      DebuggerAgent.getSink().getProbeStatusSink().addError(getProbeId(), ex);
       return false;
     } finally {
       LOGGER.debug(
@@ -199,7 +199,7 @@ public String toString() {
         language,
         location,
         probeCondition,
-        probeId,
+        getProbeId(),
         sampling,
         tagMap,
         Arrays.toString(tags),

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java

@@ -18,6 +18,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.params.provider.Arguments.arguments;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -444,15 +445,6 @@ public void multiProbeSameMethod() throws IOException, URISyntaxException {
     assertCaptureReturnValue(snapshot1.getCaptures().getReturn(), "int", "31");
   }
 
-  private List<Snapshot> assertSnapshots(
-      TestSnapshotListener listener, int expectedCount, ProbeId... probeIds) {
-    assertEquals(expectedCount, listener.snapshots.size());
-    for (int i = 0; i < probeIds.length; i++) {
-      assertEquals(probeIds[i].getId(), listener.snapshots.get(i).getProbe().getId());
-    }
-    return listener.snapshots;
-  }
-
   @Test
   public void catchBlock() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot02";
@@ -908,9 +900,14 @@ public void fieldExtractorDepth1() throws IOException, URISyntaxException {
   @Test
   public void fieldExtractorDuplicateUnionDepth() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot04";
-    LogProbe.Builder builder = createProbeBuilder(PROBE_ID, CLASS_NAME, "createSimpleData", "()");
-    LogProbe probe1 = builder.capture(0, 100, 50, Limits.DEFAULT_FIELD_COUNT).build();
-    LogProbe probe2 = builder.capture(3, 100, 50, Limits.DEFAULT_FIELD_COUNT).build();
+    LogProbe probe1 =
+        createProbeBuilder(PROBE_ID1, CLASS_NAME, "createSimpleData", "()")
+            .capture(0, 100, 50, Limits.DEFAULT_FIELD_COUNT)
+            .build();
+    LogProbe probe2 =
+        createProbeBuilder(PROBE_ID2, CLASS_NAME, "createSimpleData", "()")
+            .capture(3, 100, 50, Limits.DEFAULT_FIELD_COUNT)
+            .build();
     TestSnapshotListener listener = installProbes(probe1, probe2);
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.onClass(testClass).call("main", "").get();
@@ -2330,7 +2327,8 @@ public void enumConstructorArgs() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot23";
     final String ENUM_CLASS = CLASS_NAME + "$MyEnum";
     TestSnapshotListener listener =
-        installProbes(createMethodProbe(PROBE_ID, ENUM_CLASS, "<init>", null));
+        installProbes(
+            createProbeBuilder(PROBE_ID, ENUM_CLASS, "<init>", null).sampling(10).build());
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.onClass(testClass).call("main", "").get();
     assertEquals(2, result);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturingTestBase.java

@@ -14,6 +14,7 @@
 import com.datadog.debugger.probe.Sampled;
 import com.datadog.debugger.sink.DebuggerSink;
 import com.datadog.debugger.sink.ProbeStatusSink;
+import com.datadog.debugger.sink.Snapshot;
 import com.datadog.debugger.util.MoshiHelper;
 import com.datadog.debugger.util.MoshiSnapshotTestHelper;
 import com.datadog.debugger.util.SerializerWithLimits;
@@ -160,6 +161,15 @@ private static Collection<?> getCollection(CapturedContext.CapturedValue capture
     }
   }
 
+  protected List<Snapshot> assertSnapshots(
+      TestSnapshotListener listener, int expectedCount, ProbeId... probeIds) {
+    assertEquals(expectedCount, listener.snapshots.size());
+    for (int i = 0; i < probeIds.length; i++) {
+      assertEquals(probeIds[i].getId(), listener.snapshots.get(i).getProbe().getId());
+    }
+    return listener.snapshots;
+  }
+
   protected void assertCaptureFields(
       CapturedContext context, String name, String typeName, Map<Object, Object> expectedMap) {
     CapturedContext.CapturedValue field = getFields(context.getArguments().get("this")).get(name);