fix(sampling): always use root span for sampling decision on partial flushes (backport #14213 to 3.11) (#14227)

Ensure local root span is used to make sampling decisions when partial
flushing is enabled

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Vítor De Araújo <[email protected]>

Author: mabdinur

ddtrace/_trace/processor/__init__.py

@@ -157,7 +157,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

@@ -549,6 +549,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"