[Debugger] DEBUG-5828 Coordinate snapshot sampling per trace

tracer/src/Datadog.Trace/Debugger/Expressions/ProbeProcessor.cs

@@ -21,6 +21,7 @@
 using Datadog.Trace.Logging;
 using Datadog.Trace.SourceGenerators;
 using Datadog.Trace.Telemetry.Metrics;
+using Datadog.Trace.Util;
 using Datadog.Trace.Vendors.Serilog.Events;
 
 namespace Datadog.Trace.Debugger.Expressions
@@ -122,35 +123,84 @@ or MethodState.ExitStart
                 or MethodState.ExitStartAsync;
         }
 
+        private static bool ApplySamplingDecision(ProbeType probeType, DebuggerSamplingDecision samplingDecision)
+        {
+            switch (samplingDecision)
+            {
+                case DebuggerSamplingDecision.Keep:
+                    return true;
+                case DebuggerSamplingDecision.DropGlobal:
+                    DebuggerGuardrailMetrics.RecordEventsSkipped(probeType, MetricTags.DebuggerEventsSkippedReason.RateLimitGlobal);
+                    return false;
+                case DebuggerSamplingDecision.DropProbe:
+                    DebuggerGuardrailMetrics.RecordEventsSkipped(probeType, MetricTags.DebuggerEventsSkippedReason.RateLimitProbe);
+                    return false;
+                default:
+                    ThrowHelper.ThrowArgumentOutOfRangeException(nameof(samplingDecision), samplingDecision, null);
+                    return false;
+            }
+        }
+
         public bool TryBeginProcess(in ProbeData probeData, [NotNullWhen(true)] out IDebuggerSnapshotCreator? snapshotCreator)
         {
             var state = _state;
-            if (!state.HasCondition && !SamplePayload(state.ProbeInfo, probeData.Sampler))
+            SpanContext? activeSpanContext = null;
+            var activeSpanContextResolved = false;
+            if (!state.HasCondition && state.ShouldCoordinateSampling)
+            {
+                activeSpanContext = GetActiveScope()?.Span.Context;
+                activeSpanContextResolved = true;
+            }
+
+            if (!state.HasCondition && !SamplePayload(state, probeData.Sampler, activeSpanContext?.TraceContext))
             {
                 snapshotCreator = null;
                 return false;
             }
 
-            snapshotCreator = new DebuggerSnapshotCreator(state);
+            if (!activeSpanContextResolved &&
+                state.ProbeInfo.ProbeType is ProbeType.Log or ProbeType.Snapshot)
+            {
+                activeSpanContext = GetActiveScope()?.Span.Context;
+            }
+
+            snapshotCreator = new DebuggerSnapshotCreator(state, activeSpanContext);
             return true;
         }
 
-        private bool SamplePayload(in ProbeInfo probeInfo, IAdaptiveSampler sampler)
+        private bool SamplePayload(ProbeProcessorState state, IAdaptiveSampler sampler, TraceContext? traceContext)
+        {
+            if (!state.ShouldCoordinateSampling || traceContext is null)
+            {
+                return SamplePayloadIndependently(state, sampler);
+            }
+
+            // The first capturing probe is intentionally a trace-admission decision. Once the trace is kept,
+            // other capturing probes bypass their global and per-probe samplers and are capped once per probe.
+            var samplingDecisionProvider = new SamplingDecisionProvider(this, state, sampler);
+            var samplingDecision = traceContext.GetOrCreateDebuggerSamplingCoordinator().TrySample(state.ProbeInfo.ProbeId, samplingDecisionProvider);
+            return ApplySamplingDecision(state.ProbeInfo.ProbeType, samplingDecision);
+        }
+
+        private bool SamplePayloadIndependently(ProbeProcessorState state, IAdaptiveSampler sampler)
+        {
+            return ApplySamplingDecision(state.ProbeInfo.ProbeType, GetSamplingDecision(state, sampler));
+        }
+
+        private DebuggerSamplingDecision GetSamplingDecision(ProbeProcessorState state, IAdaptiveSampler sampler)
         {
             // Global-first matches Java; it can affect per-probe fairness and may be improved later.
-            if (probeInfo.ProbeType == ProbeType.Snapshot && !_globalRateLimiter.ShouldSampleSnapshot(probeInfo.ProbeId))
+            if (state.ShouldCoordinateSampling && !_globalRateLimiter.ShouldSampleSnapshot(state.ProbeInfo.ProbeId))
             {
-                DebuggerGuardrailMetrics.RecordEventsSkipped(probeInfo.ProbeType, MetricTags.DebuggerEventsSkippedReason.RateLimitGlobal);
-                return false;
+                return DebuggerSamplingDecision.DropGlobal;
             }
 
             if (!sampler.Sample())
             {
-                DebuggerGuardrailMetrics.RecordEventsSkipped(probeInfo.ProbeType, MetricTags.DebuggerEventsSkippedReason.RateLimitProbe);
-                return false;
+                return DebuggerSamplingDecision.DropProbe;
             }
 
-            return true;
+            return DebuggerSamplingDecision.Keep;
         }
 
         public bool Process<TCapture>(ref CaptureInfo<TCapture> info, IDebuggerSnapshotCreator inSnapshotCreator, in ProbeData probeData)
@@ -340,6 +390,12 @@ private ExpressionEvaluationResult Evaluate(ProbeProcessorState state, ProbeInfo
         {
             var evaluationResult = EvaluateCore(state, probeInfo, snapshotCreator, out shouldStopCapture, sampler);
 
+            if (evaluationResult.HasError &&
+                probeInfo.ProbeType is ProbeType.Metric or ProbeType.SpanDecoration)
+            {
+                snapshotCreator.SetActiveSpanContext(GetActiveScope()?.Span.Context);
+            }
+
             // An exceeded time budget fails open: the event is still emitted with its evaluation errors so the
             // customer can see why the probe is too slow. Only report it as skipped when the event is dropped.
             if (shouldStopCapture && evaluationResult.EvaluationBudget.TimedOut)
@@ -396,6 +452,10 @@ private ExpressionEvaluationResult EvaluateCore(ProbeProcessorState state, Probe
 
             if (captureExpressionsEvaluated && evaluationResult.IsNull())
             {
+                // The per-trace snapshot slot stays claimed even though nothing is emitted. Releasing it would let
+                // every later hit in a kept trace capture again without any sampler throttling.
+                // Not supported: if the probe is updated mid-trace to a version that captures values, it won't emit
+                // for the rest of that trace.
                 shouldStopCapture = true;
                 return evaluationResult;
             }
@@ -446,7 +506,7 @@ private ExpressionEvaluationResult EvaluateCore(ProbeProcessorState state, Probe
 
             if (evaluationResult.Condition != null && // i.e. not a metric, span probe, or span decoration
                 (evaluationResult.Condition is false ||
-                !SamplePayload(in probeInfo, sampler)))
+                !SamplePayload(state, sampler, snapshotCreator.TraceContext)))
             {
                 // if the expression evaluated to false, or there is a rate limit, stop capture
                 shouldStopCapture = true;
@@ -545,28 +605,18 @@ private bool TryGetScope(in ProbeInfo probeInfo, [NotNullWhen(true)] out Scope?
         {
             try
             {
-                if (Tracer.Instance.ActiveScope is Scope activeScope)
+                scope = GetActiveScope();
+                if (scope is not null)
                 {
-                    scope = activeScope;
                     return true;
                 }
-#if NETFRAMEWORK
-                var ctx = WcfCommon.GetCurrentOperationContext?.Invoke();
-                if (ctx?.DuckCast<IOperationContextStruct>() is { } ctxProxy
-                 && ((IDuckType?)ctxProxy.RequestContext)?.Instance is { } requestContextInstance
-                 && WcfCommon.Scopes.TryGetValue(requestContextInstance, out scope))
-                {
-                    return scope != null;
-                }
 
+#if NETFRAMEWORK
                 Log.Warning("Unable to find active scope in WCF context for span decoration. Probe ID: {ProbeId}", probeInfo.ProbeId);
-                scope = null;
-                return false;
 #else
                 Log.Warning("No active scope available for span decoration. Probe ID: {ProbeId}", probeInfo.ProbeId);
-                scope = null;
-                return false;
 #endif
+                return false;
             }
             catch (Exception e)
             {
@@ -576,6 +626,36 @@ private bool TryGetScope(in ProbeInfo probeInfo, [NotNullWhen(true)] out Scope?
             }
         }
 
+        internal static Scope? GetActiveScope()
+        {
+            if (Tracer.Instance.InternalActiveScope is { } activeScope)
+            {
+                return activeScope;
+            }
+
+#if NETFRAMEWORK
+            return GetActiveScopeFromWcf();
+#else
+            return null;
+#endif
+        }
+
+#if NETFRAMEWORK
+        private static Scope? GetActiveScopeFromWcf()
+        {
+            var operationContext = WcfCommon.GetCurrentOperationContext?.Invoke();
+            if (operationContext is not null
+             && operationContext.TryDuckCast<IOperationContextStruct>(out var operationContextProxy)
+             && operationContextProxy.RequestContext?.Instance is { } requestContextInstance
+             && WcfCommon.Scopes.TryGetValue(requestContextInstance, out var scope))
+            {
+                return scope;
+            }
+
+            return null;
+        }
+#endif
+
         internal static void AddAsyncMethodArguments<T>(DebuggerSnapshotCreator snapshotCreator, ref CaptureInfo<T> captureInfo)
         {
             var asyncCaptureInfo = captureInfo.AsyncCaptureInfo;
@@ -765,6 +845,22 @@ private void ProcessLine<TCapture>(ProbeProcessorState state, in ProbeInfo probe
             }
         }
 
+        private readonly struct SamplingDecisionProvider : IDebuggerSamplingDecisionProvider
+        {
+            private readonly ProbeProcessor _processor;
+            private readonly ProbeProcessorState _state;
+            private readonly IAdaptiveSampler _sampler;
+
+            public SamplingDecisionProvider(ProbeProcessor processor, ProbeProcessorState state, IAdaptiveSampler sampler)
+            {
+                _processor = processor;
+                _state = state;
+                _sampler = sampler;
+            }
+
+            public DebuggerSamplingDecision Sample() => _processor.GetSamplingDecision(_state, _sampler);
+        }
+
         internal sealed class ProbeProcessorState
         {
             private ProbeExpressionEvaluator? _evaluator;
@@ -789,6 +885,7 @@ private ProbeProcessorState(
                 HasCondition = condition.HasValue;
                 IsMetricCountWithoutExpression = probeInfo.ProbeType == ProbeType.Metric && (metric?.Json == null) && probeInfo.MetricKind == MetricKind.COUNT;
                 ShouldCaptureExpressions = !probeInfo.IsFullSnapshot && captureExpressions is { Length: > 0 };
+                ShouldCoordinateSampling = probeInfo.IsFullSnapshot || ShouldCaptureExpressions;
                 ShouldEvaluateExpressions =
                     HasCondition ||
                     templates is { Length: > 0 } ||
@@ -818,6 +915,8 @@ private ProbeProcessorState(
 
             internal bool ShouldCaptureExpressions { get; }
 
+            internal bool ShouldCoordinateSampling { get; }
+
             internal bool ShouldEvaluateExpressions { get; }
 
             internal static ProbeProcessorState Create(ProbeDefinition probe, int maxEvaluationTimeInMilliseconds)

tracer/src/Datadog.Trace/Debugger/RateLimiting/DebuggerSamplingCoordinator.cs

@@ -0,0 +1,59 @@
+// <copyright file="DebuggerSamplingCoordinator.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+using System.Collections.Generic;
+using System.Threading;
+
+namespace Datadog.Trace.Debugger.RateLimiting
+{
+    internal sealed class DebuggerSamplingCoordinator
+    {
+        private HashSet<string>? _emittedProbeIds;
+        private int _decision;
+
+        internal DebuggerSamplingDecision TrySample<TSamplingDecisionProvider>(string probeId, TSamplingDecisionProvider samplingDecisionProvider)
+            where TSamplingDecisionProvider : struct, IDebuggerSamplingDecisionProvider
+        {
+            var decision = (DebuggerSamplingDecision)Volatile.Read(ref _decision);
+            if (decision is DebuggerSamplingDecision.DropGlobal or DebuggerSamplingDecision.DropProbe)
+            {
+                return decision;
+            }
+
+            // The lock is reentrant. Holding it here means Sample() ran customer code on this thread
+            // (for example a first-chance exception handler) that hit another probe in this trace.
+            if (Monitor.IsEntered(this))
+            {
+                return DebuggerSamplingDecision.DropProbe;
+            }
+
+            lock (this)
+            {
+                if (_decision == (int)DebuggerSamplingDecision.Undecided)
+                {
+                    var samplingDecision = samplingDecisionProvider.Sample();
+                    if (samplingDecision == DebuggerSamplingDecision.Keep)
+                    {
+                        _emittedProbeIds = [probeId];
+                    }
+
+                    _decision = (int)samplingDecision;
+                    return samplingDecision;
+                }
+
+                if (_decision == (int)DebuggerSamplingDecision.Keep)
+                {
+                    return _emittedProbeIds!.Add(probeId)
+                               ? DebuggerSamplingDecision.Keep
+                               : DebuggerSamplingDecision.DropProbe;
+                }
+
+                return (DebuggerSamplingDecision)_decision;
+            }
+        }
+    }
+}

tracer/src/Datadog.Trace/Debugger/RateLimiting/IDebuggerSamplingDecisionProvider.cs

@@ -0,0 +1,22 @@
+// <copyright file="IDebuggerSamplingDecisionProvider.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+namespace Datadog.Trace.Debugger.RateLimiting
+{
+    internal enum DebuggerSamplingDecision
+    {
+        Undecided,
+        Keep,
+        DropGlobal,
+        DropProbe
+    }
+
+    internal interface IDebuggerSamplingDecisionProvider
+    {
+        DebuggerSamplingDecision Sample();
+    }
+}

tracer/src/Datadog.Trace/Debugger/Snapshots/DebuggerSnapshotCreator.cs

@@ -37,6 +37,7 @@ internal class DebuggerSnapshotCreator : IDebuggerSnapshotCreator, IDisposable
         private readonly CaptureLimitInfo _limitInfo;
         private readonly Func<string> _serviceNameProvider;
         private readonly Func<string?> _processTagsProvider;
+        private SpanContext? _activeSpanContext;
 
         private long _lastSampledTime;
         private TimeSpan _accumulatedDuration;
@@ -60,6 +61,11 @@ internal class DebuggerSnapshotCreator : IDebuggerSnapshotCreator, IDisposable
         private LocalsOrArgsContainer _localsOrArgsOpen;
 
         public DebuggerSnapshotCreator(bool isFullSnapshot, ProbeLocation location, bool hasCondition, string[] tags, CaptureLimitInfo limitInfo, Func<string?> processTagsProvider, Func<string> serviceNameProvider)
+            : this(isFullSnapshot, location, hasCondition, tags, limitInfo, processTagsProvider, serviceNameProvider, ProbeProcessor.GetActiveScope()?.Span.Context)
+        {
+        }
+
+        private DebuggerSnapshotCreator(bool isFullSnapshot, ProbeLocation location, bool hasCondition, string[] tags, CaptureLimitInfo limitInfo, Func<string?> processTagsProvider, Func<string> serviceNameProvider, SpanContext? activeSpanContext)
         {
             _isFullSnapshot = isFullSnapshot;
             _probeLocation = location;
@@ -75,13 +81,14 @@ public DebuggerSnapshotCreator(bool isFullSnapshot, ProbeLocation location, bool
             _limitInfo = limitInfo;
             _processTagsProvider = processTagsProvider;
             _serviceNameProvider = serviceNameProvider;
+            _activeSpanContext = activeSpanContext;
             _accumulatedDuration = new TimeSpan(0, 0, 0, 0, 0);
             _scopeMembersPool = new ObjectPool<MethodScopeMembers, MethodScopeMembersParameters>();
             Initialize();
         }
 
-        internal DebuggerSnapshotCreator(ProbeProcessor.ProbeProcessorState probeProcessorState)
-            : this(probeProcessorState.ProbeInfo.IsFullSnapshot, probeProcessorState.ProbeInfo.ProbeLocation, probeProcessorState.ProbeInfo.HasCondition, probeProcessorState.ProbeInfo.Tags, probeProcessorState.ProbeInfo.CaptureLimitInfo, DebuggerManager.ProcessTagsProvider, DebuggerManager.ServiceNameProvider)
+        internal DebuggerSnapshotCreator(ProbeProcessor.ProbeProcessorState probeProcessorState, SpanContext? activeSpanContext)
+            : this(probeProcessorState.ProbeInfo.IsFullSnapshot, probeProcessorState.ProbeInfo.ProbeLocation, probeProcessorState.ProbeInfo.HasCondition, probeProcessorState.ProbeInfo.Tags, probeProcessorState.ProbeInfo.CaptureLimitInfo, DebuggerManager.ProcessTagsProvider, DebuggerManager.ServiceNameProvider, activeSpanContext)
         {
             ProbeProcessorState = probeProcessorState;
         }
@@ -114,6 +121,8 @@ internal string SnapshotId
 
         internal ProbeProcessor.ProbeProcessorState? ProbeProcessorState { get; }
 
+        internal TraceContext? TraceContext => _activeSpanContext?.TraceContext;
+
         internal uint IncompleteReasons => _incompleteReasons;
 
         internal bool ProbeHasCondition { get; }
@@ -134,6 +143,11 @@ internal CaptureBehaviour CaptureBehaviour
             }
         }
 
+        internal void SetActiveSpanContext(SpanContext? activeSpanContext)
+        {
+            _activeSpanContext = activeSpanContext;
+        }
+
         internal void StartSampling()
         {
             _lastSampledTime = Stopwatch.GetTimestamp();
@@ -952,11 +966,9 @@ internal string FinalizeMethodSnapshot<T>(string? probeId, int probeVersion, ref
 
         internal void FinalizeSnapshot(string? methodName, string? typeFullName, string? probeFilePath)
         {
-            var activeScope = Tracer.Instance.InternalActiveScope;
-
             // TODO: support 128-bit trace ids?
-            var traceId = activeScope?.Span.TraceId128.Lower.ToString(CultureInfo.InvariantCulture);
-            var spanId = activeScope?.Span.SpanId.ToString(CultureInfo.InvariantCulture);
+            var traceId = _activeSpanContext?.TraceId128.Lower.ToString(CultureInfo.InvariantCulture);
+            var spanId = _activeSpanContext?.SpanId.ToString(CultureInfo.InvariantCulture);
 
             AddStackInfo()
             .EndSnapshot()