Add adaptive_sampling_target setting & fix logic (#1549)

hmstepanek · lrafeei · web-flow · commit 08dbcaf20b6b · 2025-10-23T12:54:00.000-07:00
* Fix removal of process attr call

* Add adaptive_sampling_target setting &amp; fix logic

* Add a new setting: `distributed_tracing.sampler.adaptive_sampling_target` that can be
used to configure the sampling target.
* Previously the logic for remote parent sampled was slightly incorrect. The sampling
value from the tracestate header was used to determine remote parent sampled, instead of
from the parent header-this has been fixed. Additionally, the sampling value was only
grabbed from the trace state and newrelic headers if there was a priority-this has been
fixed too.
* Log statements have been added to assist with sampling debug.

---------

Co-authored-by: Lalleh Rafeei &lt;84813886+lrafeei@users.noreply.github.com&gt;
diff --git a/newrelic/api/transaction.py b/newrelic/api/transaction.py
@@ -285,7 +285,8 @@ def __init__(self, application, enabled=None, source=None):
         self.tracestate = ""
         self._priority = None
         self._sampled = None
-        self._traceparent_sampled = None
+        # Remote parent sampled is set from the W3C parent header or the Newrelic header if no W3C parent header is present.
+        self._remote_parent_sampled = None
 
         self._distributed_trace_state = 0
 
@@ -569,7 +570,7 @@ def __exit__(self, exc, value, tb):
         if self._settings.distributed_tracing.enabled:
             # Sampled and priority need to be computed at the end of the
             # transaction when distributed tracing or span events are enabled.
-            self._compute_sampled_and_priority()
+            self._make_sampling_decision()
 
         self._cached_path._name = self.path
         agent_attributes = self.agent_attributes
@@ -1004,35 +1005,91 @@ def _update_agent_attributes(self):
     def user_attributes(self):
         return create_attributes(self._custom_params, DST_ALL, self.attribute_filter)
 
-    def sampling_algo_compute_sampled_and_priority(self):
-        if self._priority is None:
+    def sampling_algo_compute_sampled_and_priority(self, priority, sampled):
+        # self._priority and self._sampled are set when parsing the W3C tracestate
+        # or newrelic DT headers and may be overridden in _make_sampling_decision
+        # based on the configuration. The only time they are set in here is when the
+        # sampling decision must be made by the adaptive sampling algorithm.
+        if priority is None:
             # Truncate priority field to 6 digits past the decimal.
-            self._priority = float(f"{random.random():.6f}")  # noqa: S311
-        if self._sampled is None:
-            self._sampled = self._application.compute_sampled()
-            if self._sampled:
-                self._priority += 1
-
-    def _compute_sampled_and_priority(self):
-        if self._traceparent_sampled is None:
+            priority = float(f"{random.random():.6f}")  # noqa: S311
+        if sampled is None:
+            _logger.debug("No trusted account id found. Sampling decision will be made by adaptive sampling algorithm.")
+            sampled = self._application.compute_sampled()
+            if sampled:
+                priority += 1
+        return priority, sampled
+
+    def _compute_sampled_and_priority(
+        self,
+        priority,
+        sampled,
+        remote_parent_sampled_path,
+        remote_parent_sampled_setting,
+        remote_parent_not_sampled_path,
+        remote_parent_not_sampled_setting,
+    ):
+        if self._remote_parent_sampled is None:
             config = "default"  # Use sampling algo.
-        elif self._traceparent_sampled:
-            setting_path = "distributed_tracing.sampler.remote_parent_sampled"
-            config = self.settings.distributed_tracing.sampler.remote_parent_sampled
-        else:  # self._traceparent_sampled is False.
-            setting_path = "distributed_tracing.sampler.remote_parent_not_sampled"
-            config = self.settings.distributed_tracing.sampler.remote_parent_not_sampled
-
+            _logger.debug("Sampling decision made based on no remote parent sampling decision present.")
+        elif self._remote_parent_sampled:
+            setting_path = remote_parent_sampled_path
+            config = remote_parent_sampled_setting
+            _logger.debug(
+                "Sampling decision made based on remote_parent_sampled=%s and %s=%s.",
+                self._remote_parent_sampled,
+                setting_path,
+                config,
+            )
+        else:  # self._remote_parent_sampled is False.
+            setting_path = remote_parent_not_sampled_path
+            config = remote_parent_not_sampled_setting
+            _logger.debug(
+                "Sampling decision made based on remote_parent_sampled=%s and %s=%s.",
+                self._remote_parent_sampled,
+                setting_path,
+                config,
+            )
         if config == "always_on":
-            self._sampled = True
-            self._priority = 2.0
+            sampled = True
+            priority = 2.0
         elif config == "always_off":
-            self._sampled = False
-            self._priority = 0
+            sampled = False
+            priority = 0
         else:
-            if config != "default":
+            if config not in ("default", "adaptive"):
                 _logger.warning("%s=%s is not a recognized value. Using 'default' instead.", setting_path, config)
-            self.sampling_algo_compute_sampled_and_priority()
+
+            _logger.debug(
+                "Let adaptive sampler algorithm decide based on sampled=%s and priority=%s.", sampled, priority
+            )
+            priority, sampled = self.sampling_algo_compute_sampled_and_priority(priority, sampled)
+        return priority, sampled
+
+    def _make_sampling_decision(self):
+        # The sampling decision is computed each time a DT header is generated for exit spans as it is needed
+        # to send the DT headers. Don't recompute the sampling decision multiple times as it is expensive.
+        if hasattr(self, "_sampling_decision_made"):
+            return
+        priority = self._priority
+        sampled = self._sampled
+        _logger.debug(
+            "Full granularity tracing is enabled. Asking if full granularity wants to sample. priority=%s, sampled=%s",
+            priority,
+            sampled,
+        )
+        computed_priority, computed_sampled = self._compute_sampled_and_priority(
+            priority,
+            sampled,
+            remote_parent_sampled_path="distributed_tracing.sampler.remote_parent_sampled",
+            remote_parent_sampled_setting=self.settings.distributed_tracing.sampler.remote_parent_sampled,
+            remote_parent_not_sampled_path="distributed_tracing.sampler.remote_parent_not_sampled",
+            remote_parent_not_sampled_setting=self.settings.distributed_tracing.sampler.remote_parent_not_sampled,
+        )
+        _logger.debug("Full granularity sampling decision was %s with priority=%s.", sampled, priority)
+        self._priority = computed_priority
+        self._sampled = computed_sampled
+        self._sampling_decision_made = True
 
     def _freeze_path(self):
         if self._frozen_path is None:
@@ -1101,7 +1158,7 @@ def _create_distributed_trace_data(self):
         if not (account_id and application_id and trusted_account_key and settings.distributed_tracing.enabled):
             return
 
-        self._compute_sampled_and_priority()
+        self._make_sampling_decision()
         data = {
             "ty": "App",
             "ac": account_id,
@@ -1204,7 +1261,7 @@ def _accept_distributed_trace_payload(self, payload, transport_type="HTTP"):
             if not any(k in data for k in ("id", "tx")):
                 self._record_supportability("Supportability/DistributedTrace/AcceptPayload/ParseException")
                 return False
-
+            self._remote_parent_sampled = data.get("sa")
             settings = self._settings
             account_id = data.get("ac")
             trusted_account_key = settings.trusted_account_key or (
@@ -1254,10 +1311,8 @@ def _accept_distributed_trace_data(self, data, transport_type):
 
         self._trace_id = data.get("tr")
 
-        priority = data.get("pr")
-        if priority is not None:
-            self._priority = priority
-            self._sampled = data.get("sa")
+        self._priority = data.get("pr")
+        self._sampled = data.get("sa")
 
         if "ti" in data:
             transport_start = data["ti"] / 1000.0
@@ -1297,6 +1352,7 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
             try:
                 traceparent = ensure_str(traceparent).strip()
                 data = W3CTraceParent.decode(traceparent)
+                self._remote_parent_sampled = data.pop("sa", None)
             except:
                 data = None
 
@@ -1332,7 +1388,6 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
                     else:
                         self._record_supportability("Supportability/TraceContext/TraceState/NoNrEntry")
 
-            self._traceparent_sampled = data.get("sa")
             self._accept_distributed_trace_data(data, transport_type)
             self._record_supportability("Supportability/TraceContext/Accept/Success")
             return True
diff --git a/newrelic/config.py b/newrelic/config.py
@@ -404,6 +404,7 @@ def _process_configuration(section):
     _process_setting(section, "ml_insights_events.enabled", "getboolean", None)
     _process_setting(section, "distributed_tracing.enabled", "getboolean", None)
     _process_setting(section, "distributed_tracing.exclude_newrelic_header", "getboolean", None)
+    _process_setting(section, "distributed_tracing.sampler.adaptive_sampling_target", "getint", None)
     _process_setting(section, "distributed_tracing.sampler.remote_parent_sampled", "get", None)
     _process_setting(section, "distributed_tracing.sampler.remote_parent_not_sampled", "get", None)
     _process_setting(section, "span_events.enabled", "getboolean", None)
diff --git a/newrelic/core/agent_protocol.py b/newrelic/core/agent_protocol.py
@@ -297,6 +297,9 @@ def _connect_payload(app_name, linked_applications, environment, settings):
         connect_settings["browser_monitoring.loader"] = settings["browser_monitoring.loader"]
         connect_settings["browser_monitoring.debug"] = settings["browser_monitoring.debug"]
         connect_settings["ai_monitoring.enabled"] = settings["ai_monitoring.enabled"]
+        connect_settings["distributed_tracing.sampler.adaptive_sampling_target"] = settings[
+            "distributed_tracing.sampler.adaptive_sampling_target"
+        ]
 
         security_settings = {}
         security_settings["capture_params"] = settings["capture_params"]
diff --git a/newrelic/core/config.py b/newrelic/core/config.py
@@ -836,6 +836,9 @@ def default_otlp_host(host):
 _settings.ml_insights_events.enabled = False
 
 _settings.distributed_tracing.enabled = _environ_as_bool("NEW_RELIC_DISTRIBUTED_TRACING_ENABLED", default=True)
+_settings.distributed_tracing.sampler.adaptive_sampling_target = _environ_as_int(
+    "NEW_RELIC_DISTRIBUTED_TRACING_SAMPLER_ADAPTIVE_SAMPLING_TARGET", default=10
+)
 _settings.distributed_tracing.sampler.remote_parent_sampled = os.environ.get(
     "NEW_RELIC_DISTRIBUTED_TRACING_SAMPLER_REMOTE_PARENT_SAMPLED", "default"
 )
diff --git a/tests/agent_features/test_distributed_tracing.py b/tests/agent_features/test_distributed_tracing.py
@@ -39,6 +39,7 @@
 from newrelic.api.wsgi_application import wsgi_application
 from newrelic.core.attribute import Attribute
 
+# ruff: noqa: UP031
 distributed_trace_intrinsics = ["guid", "traceId", "priority", "sampled"]
 inbound_payload_intrinsics = [
     "parent.type",
@@ -419,24 +420,41 @@ def _test_inbound_dt_payload_acceptance():
 
 
 @pytest.mark.parametrize(
-    "sampled,remote_parent_sampled,remote_parent_not_sampled,expected_sampled,expected_priority,expected_adaptive_sampling_algo_called",
+    "traceparent_sampled,newrelic_sampled,remote_parent_sampled_setting,remote_parent_not_sampled_setting,expected_sampled,expected_priority,expected_adaptive_sampling_algo_called",
     (
-        (True, "default", "default", None, None, True),  # Uses sampling algo.
-        (True, "always_on", "default", True, 2, False),  # Always sampled.
-        (True, "always_off", "default", False, 0, False),  # Never sampled.
-        (False, "default", "default", None, None, True),  # Uses sampling algo.
-        (False, "always_on", "default", None, None, True),  # Uses sampling alog.
-        (False, "always_off", "default", None, None, True),  # Uses sampling algo.
-        (True, "default", "always_on", None, None, True),  # Uses sampling algo.
-        (True, "default", "always_off", None, None, True),  # Uses sampling algo.
-        (False, "default", "always_on", True, 2, False),  # Always sampled.
-        (False, "default", "always_off", False, 0, False),  # Never sampled.
+        (True, None, "default", "default", None, None, True),  # Uses adaptive sampling algo.
+        (True, None, "always_on", "default", True, 2, False),  # Always sampled.
+        (True, None, "always_off", "default", False, 0, False),  # Never sampled.
+        (False, None, "default", "default", None, None, True),  # Uses adaptive sampling algo.
+        (False, None, "always_on", "default", None, None, True),  # Uses adaptive sampling alog.
+        (False, None, "always_off", "default", None, None, True),  # Uses adaptive sampling algo.
+        (True, None, "default", "always_on", None, None, True),  # Uses adaptive sampling algo.
+        (True, None, "default", "always_off", None, None, True),  # Uses adaptive sampling algo.
+        (False, None, "default", "always_on", True, 2, False),  # Always sampled.
+        (False, None, "default", "always_off", False, 0, False),  # Never sampled.
+        (True, True, "default", "default", True, 1.23456, False),  # Uses sampling decision in W3C TraceState header.
+        (True, False, "default", "default", False, 1.23456, False),  # Uses sampling decision in W3C TraceState header.
+        (False, False, "default", "default", False, 1.23456, False),  # Uses sampling decision in W3C TraceState header.
+        (True, False, "always_on", "default", True, 2, False),  # Always sampled.
+        (True, True, "always_off", "default", False, 0, False),  # Never sampled.
+        (False, False, "default", "always_on", True, 2, False),  # Always sampled.
+        (False, True, "default", "always_off", False, 0, False),  # Never sampled.
+        (None, True, "default", "default", True, 0.1234, False),  # Uses sampling and priority from newrelic header.
+        (None, True, "always_on", "default", True, 2, False),  # Always sampled.
+        (None, True, "always_off", "default", False, 0, False),  # Never sampled.
+        (None, False, "default", "default", False, 0.1234, False),  # Uses sampling and priority from newrelic header.
+        (None, False, "always_on", "default", False, 0.1234, False),  # Uses sampling and priority from newrelic header.
+        (None, True, "default", "always_on", True, 0.1234, False),  # Uses sampling and priority from newrelic header.
+        (None, False, "default", "always_on", True, 2, False),  # Always sampled.
+        (None, False, "default", "always_off", False, 0, False),  # Never sampled.
+        (None, None, "default", "default", None, None, True),  # Uses adaptive sampling algo.
     ),
 )
-def test_distributed_trace_w3cparent_sampling_decision(
-    sampled,
-    remote_parent_sampled,
-    remote_parent_not_sampled,
+def test_distributed_trace_remote_parent_sampling_decision_full_granularity(
+    traceparent_sampled,
+    newrelic_sampled,
+    remote_parent_sampled_setting,
+    remote_parent_not_sampled_setting,
     expected_sampled,
     expected_priority,
     expected_adaptive_sampling_algo_called,
@@ -450,18 +468,18 @@ def test_distributed_trace_w3cparent_sampling_decision(
     test_settings = _override_settings.copy()
     test_settings.update(
         {
-            "distributed_tracing.sampler.remote_parent_sampled": remote_parent_sampled,
-            "distributed_tracing.sampler.remote_parent_not_sampled": remote_parent_not_sampled,
+            "distributed_tracing.sampler.remote_parent_sampled": remote_parent_sampled_setting,
+            "distributed_tracing.sampler.remote_parent_not_sampled": remote_parent_not_sampled_setting,
             "span_events.enabled": True,
         }
     )
     if expected_adaptive_sampling_algo_called:
         function_called_decorator = validate_function_called(
-            "newrelic.api.transaction", "Transaction.sampling_algo_compute_sampled_and_priority"
+            "newrelic.core.adaptive_sampler", "AdaptiveSampler.compute_sampled"
         )
     else:
         function_called_decorator = validate_function_not_called(
-            "newrelic.api.transaction", "Transaction.sampling_algo_compute_sampled_and_priority"
+            "newrelic.core.adaptive_sampler", "AdaptiveSampler.compute_sampled"
         )
 
     @function_called_decorator
@@ -471,10 +489,20 @@ def test_distributed_trace_w3cparent_sampling_decision(
     def _test():
         txn = current_transaction()
 
-        headers = {
-            "traceparent": f"00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-{int(sampled):02x}",
-            "tracestate": "rojo=f06a0ba902b7,congo=t61rcWkgMzE",
-        }
+        if traceparent_sampled is not None:
+            headers = {
+                "traceparent": f"00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-{int(traceparent_sampled):02x}",
+                "newrelic": '{"v":[0,1],"d":{"ty":"Mobile","ac":"123","ap":"51424","id":"5f474d64b9cc9b2a","tr":"6e2fea0b173fdad0","pr":0.1234,"sa":true,"ti":1482959525577,"tx":"27856f70d3d314b7"}}',  # This header should be ignored.
+            }
+            if newrelic_sampled is not None:
+                headers["tracestate"] = (
+                    f"1@nr=0-0-1-2827902-0af7651916cd43dd-00f067aa0ba902b7-{int(newrelic_sampled)}-1.23456-1518469636035"
+                )
+        else:
+            headers = {
+                "newrelic": '{"v":[0,1],"d":{"ty":"Mobile","ac":"1","ap":"51424","id":"00f067aa0ba902b7","tr":"0af7651916cd43dd8448eb211c80319c","pr":0.1234,"sa":%s,"ti":1482959525577,"tx":"0af7651916cd43dd"}}'
+                % (str(newrelic_sampled).lower())
+            }
         accept_distributed_trace_headers(headers)
 
     _test()
diff --git a/tests/agent_unittests/test_agent_protocol.py b/tests/agent_unittests/test_agent_protocol.py
@@ -278,10 +278,11 @@ def connect_payload_asserts(
     assert len(payload_data["security_settings"]) == 2
     assert payload_data["security_settings"]["capture_params"] == CAPTURE_PARAMS
     assert payload_data["security_settings"]["transaction_tracer"] == {"record_sql": RECORD_SQL}
-    assert len(payload_data["settings"]) == 3
+    assert len(payload_data["settings"]) == 4
     assert payload_data["settings"]["browser_monitoring.loader"] == (BROWSER_MONITORING_LOADER)
     assert payload_data["settings"]["browser_monitoring.debug"] == (BROWSER_MONITORING_DEBUG)
     assert payload_data["settings"]["ai_monitoring.enabled"] is False
+    assert payload_data["settings"]["distributed_tracing.sampler.adaptive_sampling_target"] == 10
 
     utilization_len = 5
 

Original file line number	Diff line number	Diff line change
`@@ -836,6 +836,9 @@ def default_otlp_host(host):`
`836`	`836`	`_settings.ml_insights_events.enabled = False`
`837`	`837`
`838`	`838`	`_settings.distributed_tracing.enabled = _environ_as_bool("NEW_RELIC_DISTRIBUTED_TRACING_ENABLED", default=True)`
	`839`	`+_settings.distributed_tracing.sampler.adaptive_sampling_target = _environ_as_int(`
	`840`	`+ "NEW_RELIC_DISTRIBUTED_TRACING_SAMPLER_ADAPTIVE_SAMPLING_TARGET", default=10`
	`841`	`+)`
`839`	`842`	`_settings.distributed_tracing.sampler.remote_parent_sampled = os.environ.get(`
`840`	`843`	`"NEW_RELIC_DISTRIBUTED_TRACING_SAMPLER_REMOTE_PARENT_SAMPLED", "default"`
`841`	`844`	`)`