fix(sampling): always use root span for sampling decision on partial flushes

.gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.yml

@@ -1047,7 +1047,7 @@ experiments:
               - max_rss_usage < 34.00 MB
           - name: span-start-finish-traceid128
             thresholds:
-              - execution_time < 58.05 ms
+              - execution_time < 60.05 ms
               - max_rss_usage < 34.00 MB
           - name: span-start-traceid128
             thresholds:

ddtrace/_trace/processor/__init__.py

@@ -158,7 +158,7 @@ def process_trace(self, trace: List[Span]) -> Optional[List[Span]]:
 
             # only trace sample if we haven't already sampled
             if root_ctx and root_ctx.sampling_priority is None:
-                self.sampler.sample(trace[0])
+                self.sampler.sample(trace[0]._local_root)
             # When stats computation is enabled in the tracer then we can
             # safely drop the traces.
             if self._compute_stats_enabled and not self.apm_opt_out:

releasenotes/notes/partial-flush-sampling-fix-e34a7c2dd47ae82a.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    tracing: Fix inconsistent trace sampling during partial flush (traces >300 spans). Now correctly applies sampling rules to the root span instead of a random payload span. 
+    Since traces are sampled only once, using the wrong span could bypass sampling rules and cause the agent to apply default rate limits. Fixes regression introduced in v2.8.0.

tests/tracer/test_sampler.py

@@ -548,6 +548,38 @@ def test_rate_limit_without_sampling_rules_warning():
     assert config._trace_rate_limit == 2
 
 
+@pytest.mark.subprocess(
+    env={
+        "DD_TRACE_PARTIAL_FLUSH_ENABLED": "true",
+        "DD_TRACE_PARTIAL_FLUSH_MIN_SPANS": "5",
+        "DD_TRACE_SAMPLING_RULES": '[{"sample_rate":0, "name":"root_span"}, {"sample_rate":1, "name":"child_span1"}]',
+    },
+)
+def test_partial_flush_with_sampling_rules():
+    """
+    Detects a bug where the local root span is not used to make sampling decisions when a trace is partially flushed.
+    """
+    from ddtrace import tracer
+
+    with tracer.trace("root_span") as root_span:
+        with tracer.trace("child_span1") as child_span1:
+            for _ in range(4):
+                with tracer.trace("span"):
+                    pass
+
+        with tracer.trace("child_span2") as child_span2:
+            for _ in range(4):
+                with tracer.trace("span"):
+                    pass
+
+    assert root_span.get_metric("_dd.rule_psr") == 0, repr(root_span)
+    assert child_span1.get_metric("_dd.py.partial_flush") == 5, repr(child_span1)
+    assert child_span2.get_metric("_dd.py.partial_flush") == 5, repr(child_span2)
+
+    for span in (root_span, child_span1, child_span2):
+        assert span.context.sampling_priority == -1, repr(span)
+
+
 def test_datadog_sampler_init():
     sampler = DatadogSampler()
     assert sampler.rules == [], "DatadogSampler initialized with no arguments should hold no rules"