Fix stats computation pipeline

Author: andrewlock

tracer/src/Datadog.Trace/Agent/AgentWriter.cs

@@ -433,34 +433,36 @@ private void SerializeTrace(in SpanCollection spans)
             var chunk = spans;
             if (CanComputeStats)
             {
-                chunk = _statsAggregator.ProcessTrace(in chunk);
-                bool shouldSendTrace = _statsAggregator.ShouldKeepTrace(in chunk);
+                var dropReason = _statsAggregator.ProcessTrace(ref chunk);
                 _statsAggregator.AddRange(in chunk);
-                var singleSpanSamplingSpans = new List<Span>(); // TODO maybe we can store this from above?
+                List<Span>? singleSpanSamplingSpans = null; // TODO maybe we can store this from above?
 
                 foreach (var span in chunk)
                 {
                     if (span.GetMetric(Metrics.SingleSpanSampling.SamplingMechanism) is not null)
                     {
+                        singleSpanSamplingSpans ??= new();
                         singleSpanSamplingSpans.Add(span);
                     }
                 }
 
-                if (shouldSendTrace)
+                if (dropReason is null)
                 {
                     TelemetryFactory.Metrics.RecordCountTraceChunkEnqueued(MetricTags.TraceChunkEnqueueReason.P0Keep);
                     TelemetryFactory.Metrics.RecordCountSpanEnqueuedForSerialization(MetricTags.SpanEnqueueReason.P0Keep, chunk.Count);
                 }
                 else
                 {
-                    // If stats computation determined that we can drop the P0 Trace,
+                    // If stats computation determined that we can drop the trace,
                     // skip all other processing
-                    TelemetryFactory.Metrics.RecordCountTraceChunkDropped(MetricTags.DropReason.P0Drop);
-                    if (singleSpanSamplingSpans.Count == 0)
+                    var reasonTag = dropReason.Value.ToTagReason();
+
+                    TelemetryFactory.Metrics.RecordCountTraceChunkDropped(reasonTag);
+                    if (singleSpanSamplingSpans is null)
                     {
                         Interlocked.Increment(ref _droppedP0Traces);
                         Interlocked.Add(ref _droppedP0Spans, chunk.Count);
-                        TelemetryFactory.Metrics.RecordCountSpanDropped(MetricTags.DropReason.P0Drop, chunk.Count);
+                        TelemetryFactory.Metrics.RecordCountSpanDropped(reasonTag, chunk.Count);
                         return;
                     }
                     else
@@ -472,7 +474,7 @@ private void SerializeTrace(in SpanCollection spans)
                         var spansDropped = chunk.Count - singleSpanSamplingSpans.Count;
                         Interlocked.Add(ref _droppedP0Spans, spansDropped);
                         chunk = new SpanCollection(singleSpanSamplingSpans.ToArray(), singleSpanSamplingSpans.Count);
-                        TelemetryFactory.Metrics.RecordCountSpanDropped(MetricTags.DropReason.P0Drop, spansDropped);
+                        TelemetryFactory.Metrics.RecordCountSpanDropped(reasonTag, spansDropped);
                         TelemetryFactory.Metrics.RecordCountSpanEnqueuedForSerialization(MetricTags.SpanEnqueueReason.SingleSpanSampling, chunk.Count);
                         TelemetryFactory.Metrics.RecordCountTracePartialFlush(MetricTags.PartialFlushReason.SingleSpanIngestion);
                     }

tracer/src/Datadog.Trace/Agent/IStatsAggregator.cs

@@ -42,13 +42,12 @@ internal interface IStatsAggregator
         void AddRange(in SpanCollection spans);
 
         /// <summary>
-        /// Runs a series of samplers over the entire trace chunk
+        /// Apply normalization, filtering, obfuscation, and sampling, to understand if the
+        /// trace should be kept
         /// </summary>
-        /// <param name="spans">The trace chunk to sample</param>
-        /// <returns>True if the trace chunk should be sampled, false otherwise.</returns>
-        bool ShouldKeepTrace(in SpanCollection spans);
-
-        SpanCollection ProcessTrace(in SpanCollection trace);
+        /// <param name="spans">The spans chunk to process</param>
+        /// <returns>An optional trace drop reason, or <c>null</c> if the trace should _not_ be dropped</returns>
+        TraceDropReason? ProcessTrace(ref SpanCollection spans);
 
         Task DisposeAsync();
 

tracer/src/Datadog.Trace/Agent/NullStatsAggregator.cs

@@ -21,9 +21,7 @@ public void AddRange(in SpanCollection spans)
         {
         }
 
-        public bool ShouldKeepTrace(in SpanCollection spans) => true;
-
-        public SpanCollection ProcessTrace(in SpanCollection trace) => trace;
+        public TraceDropReason? ProcessTrace(ref SpanCollection spans) => null;
 
         public Task DisposeAsync()
         {

tracer/src/Datadog.Trace/Agent/StatsAggregator.cs

@@ -134,13 +134,6 @@ public void Add(params Span[] spans)
 
         public void AddRange(in SpanCollection spans)
         {
-            // Trace-level filters from the agent must be applied before stats computation.
-            // Rejected traces should not contribute to stats.
-            if (IsTraceFiltered(in spans))
-            {
-                return;
-            }
-
             // Contention around this lock is expected to be very small:
             // AddRange is called from the serialization thread, and concurrent serialization
             // of traces is a rare corner-case (happening only during shutdown).
@@ -154,21 +147,72 @@ public void AddRange(in SpanCollection spans)
             }
         }
 
-        public bool ShouldKeepTrace(in SpanCollection trace)
+        public TraceDropReason? ProcessTrace(ref SpanCollection spans)
         {
-            // For OTLP, align with the OpenTelemetry SDK behavior to export a trace based
-            // solely on its sampling decision.
-            if (_isOtlp)
+            // Follow the same processing steps as the Go tracer
+            spans = NormalizeTrace(in spans);
+            if (ShouldFilterTrace(in spans))
             {
-                return _prioritySampler.Sample(in trace);
+                return TraceDropReason.TraceFilter;
+            }
+
+            spans = ObfuscateTrace(in spans);
+            if (!ShouldKeepTrace(in spans))
+            {
+                return TraceDropReason.Unsampled;
+            }
+
+            return null; // keep
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal SpanCollection NormalizeTrace(in SpanCollection trace)
+        {
+            try
+            {
+                return _normalizerProcessor.Process(in trace);
             }
+            catch (Exception e)
+            {
+                Log.Error(e, "Error executing normalizer trace processor");
+                return trace;
+            }
+        }
 
-            // Trace-level filters from the agent must reject traces before sampling.
-            if (IsTraceFiltered(in trace))
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal bool ShouldFilterTrace(in SpanCollection trace)
+        {
+            var filter = Volatile.Read(ref _traceFilter);
+            if (filter is null)
             {
                 return false;
             }
 
+            // Find the local root span, searching from the last span, as that's where it normally is
+            var traceContext = TraceContext.GetTraceContext(in trace);
+            var localRoot = traceContext?.RootSpan;
+
+            if (localRoot is not null && trace.ContainsSpanId(localRoot.SpanId, trace.Count - 1))
+            {
+                // localRoot is in the trace chunk, so we can apply the filter directly
+                return filter.ShouldKeepTrace(localRoot);
+            }
+
+            // local root isn't in the trace chunk (can happen with partial flushing)
+            // or we don't have a local root (I don't know when that happens!)
+            return false;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal bool ShouldKeepTrace(in SpanCollection trace)
+        {
+            // For OTLP, align with the OpenTelemetry SDK behavior to export a trace based
+            // solely on its sampling decision.
+            if (_isOtlp)
+            {
+                return _prioritySampler.Sample(in trace);
+            }
+
             // Note: The RareSampler must be run before all other samplers so that
             // the first rare span in the trace chunk (if any) is marked with "_dd.rare".
             // The sampling decision is only used if no other samplers choose to keep the trace chunk.
@@ -181,33 +225,23 @@ public bool ShouldKeepTrace(in SpanCollection trace)
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public SpanCollection ProcessTrace(in SpanCollection trace)
+        internal SpanCollection ObfuscateTrace(in SpanCollection trace)
         {
-            var spans = trace;
-            try
-            {
-                spans = _normalizerProcessor.Process(in spans);
-            }
-            catch (Exception e)
-            {
-                Log.Error(e, "Error executing normalizer trace processor");
-            }
-
             // Only obfuscate resources when the tracer has negotiated obfuscation responsibility.
             // The tracer currently only supports obfuscation version 1 (SQL, Cassandra, Redis).
             if (Volatile.Read(ref _tracerObfuscationVersion) == 1)
             {
                 try
                 {
-                    spans = _obfuscatorProcessor.Process(in spans);
+                    return _obfuscatorProcessor.Process(in trace);
                 }
                 catch (Exception e)
                 {
                     Log.Error(e, "Error executing obfuscator trace processor");
                 }
             }
 
-            return spans;
+            return trace;
         }
 
         public StatsAggregationKey BuildKey(Span span, out List<byte[]> utf8PeerTags)

tracer/src/Datadog.Trace/Agent/TraceDropReason.cs

@@ -0,0 +1,26 @@
+// <copyright file="TraceDropReason.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 Datadog.Trace.Telemetry.Metrics;
+
+namespace Datadog.Trace.Agent;
+
+internal enum TraceDropReason
+{
+    TraceFilter,
+    Unsampled,
+}
+
+#pragma warning disable SA1649 //  File name must match first type name
+internal static class TraceDropReasonExtensions
+{
+    public static MetricTags.DropReason ToTagReason(this TraceDropReason reason) => reason switch
+    {
+        TraceDropReason.TraceFilter => MetricTags.DropReason.TraceFilter,
+        _ => MetricTags.DropReason.P0Drop,
+    };
+}

tracer/test/Datadog.Trace.Tests/Agent/AgentWriterTests.cs

@@ -601,9 +601,11 @@ public void AddRange(in SpanCollection spans)
                 AddedSpans.Add(spans);
             }
 
-            public bool ShouldKeepTrace(in SpanCollection spans) => shouldKeepTrace;
-
-            public SpanCollection ProcessTrace(in SpanCollection trace) => processTrace(trace);
+            public TraceDropReason? ProcessTrace(ref SpanCollection spans)
+            {
+                spans = processTrace(spans);
+                return shouldKeepTrace ? null : TraceDropReason.Unsampled;
+            }
 
             public Task DisposeAsync() => Task.CompletedTask;
 

tracer/test/Datadog.Trace.Tests/Agent/StatsAggregatorTests.cs

@@ -499,7 +499,7 @@ public void CreateStatsAggregator_Otlp_AlwaysComputesStats()
         }
 
         [Fact]
-        public async Task Otlp_ShouldKeepTraces_TrueWhenTraceSampled()
+        public async Task Otlp_ProcessTrace_WhenTraceSampled()
         {
             var aggregator = StatsAggregator.Create(Mock.Of<IApi>(), GetSettings(), NullDiscoveryService.Instance, isOtlp: true);
             await using var tracer = TracerHelper.CreateWithFakeAgent();
@@ -511,11 +511,12 @@ public async Task Otlp_ShouldKeepTraces_TrueWhenTraceSampled()
             traceContext.SetSamplingPriority(priority: SamplingPriorityValues.AutoKeep, mechanism: SamplingMechanism.LocalTraceSamplingRule, rate: null, limiterRate: null);
 
             var traceChunk = new SpanCollection([span]);
-            aggregator.ShouldKeepTrace(traceChunk).Should().BeTrue();
+            var dropReason = aggregator.ProcessTrace(ref traceChunk);
+            dropReason.Should().BeNull("sampled trace should be kept");
         }
 
         [Fact]
-        public async Task Otlp_ShouldKeepTraces_FalseWhenTraceNotSampled()
+        public async Task Otlp_ProcessTrace_WhenTraceNotSampled()
         {
             var aggregator = StatsAggregator.Create(Mock.Of<IApi>(), GetSettings(), NullDiscoveryService.Instance, isOtlp: true);
             await using var tracer = TracerHelper.CreateWithFakeAgent();
@@ -527,7 +528,8 @@ public async Task Otlp_ShouldKeepTraces_FalseWhenTraceNotSampled()
             traceContext.SetSamplingPriority(priority: SamplingPriorityValues.AutoReject, mechanism: SamplingMechanism.LocalTraceSamplingRule, rate: null, limiterRate: null);
 
             var traceChunk = new SpanCollection([span]);
-            aggregator.ShouldKeepTrace(traceChunk).Should().BeFalse();
+            var dropReason = aggregator.ProcessTrace(ref traceChunk);
+            dropReason.Should().Be(TraceDropReason.Unsampled);
         }
 
         [Fact]