Add traceparent logic

Author: hmstepanek

newrelic/api/transaction.py

@@ -286,6 +286,7 @@ def __init__(self, application, enabled=None, source=None):
         self.tracestate = ""
         self._priority = None
         self._sampled = None
+        self._traceparent_sampled = None
 
         self._distributed_trace_state = 0
 
@@ -1004,16 +1005,32 @@ def _update_agent_attributes(self):
     def user_attributes(self):
         return create_attributes(self._custom_params, DST_ALL, self.attribute_filter)
 
-    def _compute_sampled_and_priority(self):
+    def sampling_algo_compute_sampled_and_priority(self):
         if self._priority is None:
-            # truncate priority field to 6 digits past the decimal
+            # 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:
+            config = "default"  # Use sampling algo.
+        elif self._traceparent_sampled:
+            config = self.settings.distributed_tracing.sampler.remote_parent_sampled
+        else:  # self._traceparent_sampled is False.
+            config = self.settings.distributed_tracing.sampler.remote_parent_not_sampled
+
+        if config == 'always_on':
+            self._sampled = True
+            self._priority = 2.0
+        elif config == 'always_off':
+            self._sampled = False
+            self._priority = 0
+        else:
+            self.sampling_algo_compute_sampled_and_priority()
+
     def _freeze_path(self):
         if self._frozen_path is None:
             self._name_priority = None
@@ -1348,6 +1365,7 @@ 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

newrelic/common/encoding_utils.py

@@ -34,6 +34,7 @@
 DELIMITER_FORMAT_RE = re.compile("[ \t]*,[ \t]*")
 PARENT_TYPE = {"0": "App", "1": "Browser", "2": "Mobile"}
 BASE64_DECODE_STR = getattr(base64, "decodestring", None)
+FLAG_SAMPLED = 1
 
 
 # Functions for encoding/decoding JSON. These wrappers are used in order
@@ -455,7 +456,10 @@ def decode(cls, payload):
         if parent_id == "0" * 16 or trace_id == "0" * 32:
             return None
 
-        return cls(tr=trace_id, id=parent_id)
+        # Sampled flag
+        sa = bool(int(fields[3], 2) & FLAG_SAMPLED)
+
+        return cls(tr=trace_id, id=parent_id, sa=sa)
 
 
 class W3CTraceState(OrderedDict):

tests/agent_features/test_distributed_tracing.py

@@ -17,11 +17,12 @@
 
 import pytest
 import webtest
-from testing_support.fixtures import override_application_settings, validate_attributes
+from testing_support.fixtures import override_application_settings, validate_attributes, validate_attributes_complete
 from testing_support.validators.validate_error_event_attributes import validate_error_event_attributes
 from testing_support.validators.validate_transaction_event_attributes import validate_transaction_event_attributes
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 
+from newrelic.common.object_wrapper import wrap_function_wrapper
 from newrelic.api.application import application_instance
 from newrelic.api.background_task import BackgroundTask, background_task
 from newrelic.api.external_trace import ExternalTrace
@@ -36,6 +37,7 @@
 )
 from newrelic.api.web_transaction import WSGIWebTransaction
 from newrelic.api.wsgi_application import wsgi_application
+from newrelic.core.attribute import Attribute
 
 distributed_trace_intrinsics = ["guid", "traceId", "priority", "sampled"]
 inbound_payload_intrinsics = [
@@ -409,3 +411,65 @@ def _test_inbound_dt_payload_acceptance():
             assert not result
 
     _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",
+    (
+        (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.
+    )
+)
+def test_distributed_trace_w3cparent_sampling_decision(sampled, remote_parent_sampled, remote_parent_not_sampled, expected_sampled, expected_priority, expected_adaptive_sampling_algo_called, wrap_sampling_algo_compute_sampled_and_priority):
+    required_intrinsics = []
+    if expected_sampled is not None:
+        required_intrinsics.append(Attribute(name="sampled", value=expected_sampled, destinations=0b110))
+    if expected_priority is not None:
+        required_intrinsics.append(Attribute(name="priority", value=expected_priority, destinations=0b110))
+
+    wrap_function_wrapper("newrelic.api.transaction", "Transaction.sampling_algo_compute_sampled_and_priority", wrap_sampling_algo_compute_sampled_and_priority)
+
+    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,
+        "span_events.enabled": True,
+    })
+
+    @override_application_settings(test_settings)
+    @validate_attributes_complete("intrinsic", required_intrinsics)
+    @background_task(name="test_distributed_trace_attributes")
+    def _test():
+        txn = current_transaction()
+
+        headers = {
+            "traceparent": f"00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-{int(sampled):02x}",
+            "tracestate": f"rojo=f06a0ba902b7,congo=t61rcWkgMzE",
+        }
+        accept_distributed_trace_headers(headers)
+
+    _test()
+
+    assert was_called  == expected_adaptive_sampling_algo_called
+
+
+@pytest.fixture
+def wrap_sampling_algo_compute_sampled_and_priority():
+    global was_called
+    was_called = False
+
+    def _wrap_sampling_algo_compute_sampled_and_priority(wrapped, instance, args, kwargs):
+        global was_called
+        was_called = True
+        return wrapped(*args, **kwargs)
+
+    return _wrap_sampling_algo_compute_sampled_and_priority
+