fix(sampling): avoid resetting sampling rules on forks [backport 3.14] (#14602)

Backport 1692801 from #14557 to 3.14.

Change logic to avoid recreating DatadogSampler unnecessarily when the
SpanAggregator is recreated. If apm_opt_out is enabled, update the
sampler to use a new `RateLimiter` with updated values.
This was potentially causing issues with sampling rules getting reset
when the DatadogSampler gets recreated.

## 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: Quinna Halim <[email protected]>

Author: dd-octo-sts[bot]

ddtrace/_trace/processor/__init__.py

@@ -3,15 +3,12 @@
 from itertools import chain
 import logging
 from threading import RLock
-from typing import Any
 from typing import DefaultDict
 from typing import Dict
 from typing import List
 from typing import Optional
-from typing import Union
 
 from ddtrace._trace.sampler import DatadogSampler
-from ddtrace._trace.sampler import RateSampler
 from ddtrace._trace.span import Span
 from ddtrace._trace.span import _get_64_highest_order_bits_as_hex
 from ddtrace.constants import _APM_ENABLED_METRIC_KEY as MK_APM_ENABLED
@@ -22,6 +19,7 @@
 from ddtrace.internal.constants import LAST_DD_PARENT_ID_KEY
 from ddtrace.internal.constants import MAX_UINT_64BITS
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.rate_limiter import RateLimiter
 from ddtrace.internal.sampling import SpanSamplingRule
 from ddtrace.internal.sampling import get_span_sampling_rules
 from ddtrace.internal.service import ServiceStatusError
@@ -119,28 +117,30 @@ def __init__(
         compute_stats_enabled: bool,
         single_span_rules: List[SpanSamplingRule],
         apm_opt_out: bool,
-        agent_based_samplers: Optional[dict] = None,
     ):
         super(TraceSamplingProcessor, self).__init__()
         self._compute_stats_enabled = compute_stats_enabled
         self.single_span_rules = single_span_rules
+        self.sampler = DatadogSampler()
         self.apm_opt_out = apm_opt_out
 
+    @property
+    def apm_opt_out(self):
+        return self._apm_opt_out
+
+    @apm_opt_out.setter
+    def apm_opt_out(self, value):
         # If ASM is enabled but tracing is disabled,
         # we need to set the rate limiting to 1 trace per minute
         # for the backend to consider the service as alive.
-        sampler_kwargs: Dict[str, Any] = {
-            "agent_based_samplers": agent_based_samplers,
-        }
-        if self.apm_opt_out:
-            sampler_kwargs.update(
-                {
-                    "rate_limit": 1,
-                    "rate_limit_window": 60e9,
-                    "rate_limit_always_on": True,
-                }
-            )
-        self.sampler: Union[DatadogSampler, RateSampler] = DatadogSampler(**sampler_kwargs)
+        if value:
+            self.sampler.limiter = RateLimiter(rate_limit=1, time_window=60e9)
+            self.sampler._rate_limit_always_on = True
+            log.debug("Enabling apm opt out on DatadogSampler: %s", self.sampler)
+        else:
+            self.sampler.limiter = RateLimiter(rate_limit=int(config._trace_rate_limit), time_window=1e9)
+            self.sampler._rate_limit_always_on = False
+        self._apm_opt_out = value
 
     def process_trace(self, trace: List[Span]) -> Optional[List[Span]]:
         if trace:
@@ -325,7 +325,7 @@ def on_span_start(self, span: Span) -> None:
         log.debug(self.SPAN_START_DEBUG_MESSAGE, span, len(trace.spans))
 
     def on_span_finish(self, span: Span) -> None:
-        # Aqcuire lock to get finished and update trace.spans
+        # Acquire lock to get finished and update trace.spans
         with self._lock:
             integration_name = span._meta.get(COMPONENT, span._span_api)
             self._span_metrics["spans_finished"][integration_name] += 1
@@ -470,22 +470,12 @@ def reset(
         # Re-create the writer to ensure it is consistent with updated configurations (ex: api_version)
         self.writer = self.writer.recreate(appsec_enabled=appsec_enabled)
 
-        # Recreate the sampling processor using new or existing config values.
-        # If an argument is None, the current value is preserved.
-        if compute_stats is None:
-            compute_stats = self.sampling_processor._compute_stats_enabled
-        if apm_opt_out is None:
-            apm_opt_out = self.sampling_processor.apm_opt_out
-        self.sampling_processor = TraceSamplingProcessor(
-            compute_stats,
-            get_span_sampling_rules(),
-            apm_opt_out,
-            self.sampling_processor.sampler._agent_based_samplers
-            if isinstance(self.sampling_processor.sampler, DatadogSampler)
-            else None,
-        )
+        if compute_stats is not None:
+            self.sampling_processor._compute_stats_enabled = compute_stats
+
+        if apm_opt_out is not None:
+            self.sampling_processor.apm_opt_out = apm_opt_out
 
-        # Update user processors if provided.
         if user_processors is not None:
             self.user_processors = user_processors
 

ddtrace/_trace/sampler.py

@@ -91,7 +91,6 @@ def __init__(
         rate_limit: Optional[int] = None,
         rate_limit_window: float = 1e9,
         rate_limit_always_on: bool = False,
-        agent_based_samplers: Optional[Dict[str, RateSampler]] = None,
     ):
         """
         Constructor for DatadogSampler sampler
@@ -102,8 +101,6 @@ def __init__(
         :param rate_limit_window: The time window in nanoseconds for the rate limit, default is 1 second
         :param rate_limit_always_on: If set to `True`, the rate limit is always applied, even if no sampling rules
             are provided.
-        :param agent_based_samplers: A dictionary of service-based samplers, mapping a key in the format
-            `service:<service>,env:<env>` to a :class:`RateSampler` instance.
         """
         # Set sampling rules
         global_sampling_rules = config._trace_sampling_rules
@@ -112,7 +109,7 @@ def __init__(
         else:
             self.rules: List[SamplingRule] = rules or []
         # Set Agent based samplers
-        self._agent_based_samplers = agent_based_samplers or {}
+        self._agent_based_samplers: Dict = {}
         # Set rate limiter
         self._rate_limit_always_on: bool = rate_limit_always_on
         if rate_limit is None:

releasenotes/notes/fix-datadog-sampler-updates-8bfaac34f9a1598b.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    sampling: This change prevents the DatadogSampler from getting recreated whenever the SpanAggregator is reset, and instead updates the rate limiter that the sampler uses.

tests/tracer/test_processors.py

@@ -107,7 +107,7 @@ def process_trace(self, trace):
 
 def test_aggregator_reset_default_args():
     """
-    Test that on reset, the aggregator recreates the sampling processor and trace writer.
+    Test that on reset, the aggregator recreates trace writer but not the sampling processor (by default).
     Processors and trace buffers should be reset not reset.
     """
     dd_proc = DummyProcessor()
@@ -134,11 +134,46 @@ def test_aggregator_reset_default_args():
     assert dd_proc in aggr.dd_processors
     assert user_proc in aggr.user_processors
     assert aggr.writer is not dm_writer
-    assert sampling_proc is not aggr.sampling_processor
+    assert sampling_proc is aggr.sampling_processor
     assert not aggr._traces
     assert len(aggr._span_metrics["spans_created"]) == 0
 
 
+def test_aggregator_reset_apm_opt_out_preserves_sampling():
+    """
+    Test that calling aggr.reset(apm_opt_out=True) updates the apm_opt_out setting
+    but preserves the sampling rules on the TraceSamplingProcessor.
+    """
+    sampling_rule = SpanSamplingRule(service="test_service", name="test_name", sample_rate=0.5, max_per_second=10)
+
+    dd_proc = DummyProcessor()
+    user_proc = DummyProcessor()
+    aggr = SpanAggregator(
+        partial_flush_enabled=False,
+        partial_flush_min_spans=1,
+        dd_processors=[dd_proc],
+        user_processors=[user_proc],
+    )
+
+    sampling_proc = aggr.sampling_processor
+    original_apm_opt_out = sampling_proc.apm_opt_out
+
+    sampling_proc.sampler.rules = [TraceSamplingRule(sample_rate=0.1)]
+    rule = sampling_proc.sampler.rules[0]
+    sampling_proc.single_span_rules = [sampling_rule]
+
+    assert sampling_proc.single_span_rules == [sampling_rule]
+    assert sampling_proc.apm_opt_out == original_apm_opt_out
+
+    aggr.reset(apm_opt_out=True)
+
+    # Assert that sampling rules are preserved after reset
+    assert sampling_proc.apm_opt_out is True
+    assert sampling_proc.single_span_rules == [sampling_rule]
+    assert sampling_proc.sampler.rules[0] == rule
+    assert sampling_proc is aggr.sampling_processor
+
+
 @pytest.mark.parametrize("writer_class", (AgentWriter, NativeWriter))
 def test_aggregator_reset_with_args(writer_class):
     """