ref: Add sampling to span first (11)

sentry_sdk/scope.py

@@ -30,6 +30,7 @@
     Baggage,
     has_tracing_enabled,
     has_span_streaming_enabled,
+    _make_sampling_decision,
     normalize_incoming_data,
     PropagationContext,
 )
@@ -1199,6 +1200,21 @@
         if parent_span is None:
             propagation_context = self.get_active_propagation_context()
 
+            sampled, sample_rate, sample_rand, outcome = _make_sampling_decision(
+                name,
+                attributes,
+                self,
+            )
+
+            if sample_rate is not None:
+                self._update_sample_rate(sample_rate)
+
+            if sampled is False:
+                return NoOpStreamedSpan(
+                    scope=self,
+                    unsampled_reason=outcome,
+                )
+
             return StreamedSpan(
                 name=name,
                 attributes=attributes,
@@ -1209,12 +1225,14 @@
                 parent_span_id=propagation_context.parent_span_id,
                 parent_sampled=propagation_context.parent_sampled,
                 baggage=propagation_context.baggage,
+                sample_rand=sample_rand,
+                sample_rate=sample_rate,
             )
 
         # This is a child span; take propagation context from the parent span
         with new_scope():
             if isinstance(parent_span, NoOpStreamedSpan):
-                return NoOpStreamedSpan()
+                return NoOpStreamedSpan(unsampled_reason=parent_span._unsampled_reason)
 
             return StreamedSpan(
                 name=name,
@@ -1227,6 +1245,15 @@
                 parent_sampled=parent_span.sampled,
             )
 
+    def _update_sample_rate(self, sample_rate: float) -> None:
+        # If we had to adjust the sample rate when setting the sampling decision
+        # for a span, it needs to be updated in the propagation context too
+        propagation_context = self.get_active_propagation_context()
+        baggage = propagation_context.baggage
+
+        if baggage is not None:
+            baggage.sentry_items["sample_rate"] = str(sample_rate)
+
     def continue_trace(
         self,
         environ_or_headers: "Dict[str, Any]",

sentry_sdk/traces.py

@@ -224,6 +224,8 @@
         "_scope",
         "_previous_span_on_scope",
         "_baggage",
+        "_sample_rand",
+        "_sample_rate",
     )
 
     def __init__(
@@ -238,6 +240,8 @@
         parent_span_id: "Optional[str]" = None,
         parent_sampled: "Optional[bool]" = None,
         baggage: "Optional[Baggage]" = None,
+        sample_rate: "Optional[float]" = None,
+        sample_rand: "Optional[float]" = None,
     ):
         self._name: str = name
         self._active: bool = active
@@ -254,6 +258,8 @@
         self._parent_span_id = parent_span_id
         self._parent_sampled = parent_sampled
         self._baggage = baggage
+        self._sample_rand = sample_rand
+        self._sample_rate = sample_rate
 
         self._start_timestamp = datetime.now(timezone.utc)
         self._timestamp: "Optional[datetime]" = None
@@ -441,13 +447,18 @@
 
 
 class NoOpStreamedSpan(StreamedSpan):
-    __slots__ = ()
+    __slots__ = (
+        "_finished",
+        "_unsampled_reason",
+    )
 
     def __init__(
         self,
+        unsampled_reason: "Optional[str]" = None,
         scope: "Optional[sentry_sdk.Scope]" = None,
     ) -> None:
         self._scope = scope  # type: ignore[assignment]
+        self._unsampled_reason = unsampled_reason
 
         self._start()
 
@@ -471,16 +482,25 @@
         self._previous_span_on_scope = old_span
 
     def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
-        if self._scope is None:
+        if self._finished:
             return
 
-        if not hasattr(self, "_previous_span_on_scope"):
-            return
+        client = sentry_sdk.get_client()
+        if client.is_active() and client.transport:
+            logger.debug("Discarding span because sampled = False")
+            client.transport.record_lost_event(
+                reason=self._unsampled_reason or "sample_rate",
+                data_category="span",
+                quantity=1,
+            )
+
+        if self._scope:
+            with capture_internal_exceptions():
+                old_span = self._previous_span_on_scope
+                del self._previous_span_on_scope
+                self._scope.span = old_span
 
-        with capture_internal_exceptions():
-            old_span = self._previous_span_on_scope
-            del self._previous_span_on_scope
-            self._scope.span = old_span
+        self._finished = True
 
     def end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
         self._end()

sentry_sdk/tracing_utils.py

@@ -23,6 +23,7 @@
     to_string,
     try_convert,
     is_sentry_url,
+    is_valid_sample_rate,
     _is_external_source,
     _is_in_project_root,
     _module_in_list,
@@ -41,6 +42,8 @@
 
     from types import FrameType
 
+    from sentry_sdk._types import Attributes
+
 
 SENTRY_TRACE_REGEX = re.compile(
     "^[ \t]*"  # whitespace
@@ -1379,6 +1382,93 @@ def add_sentry_baggage_to_headers(
     )
 
 
+def _make_sampling_decision(
+    name: str,
+    attributes: "Optional[Attributes]",
+    scope: "sentry_sdk.Scope",
+) -> "tuple[bool, Optional[float], Optional[float], Optional[str]]":
+    """
+    Decide whether a span should be sampled.
+
+    Returns a tuple with:
+    1. the sampling decision
+    2. the effective sample rate
+    3. the sample rand
+    4. the reason for not sampling the span, if unsampled
+    """
+    client = sentry_sdk.get_client()
+
+    if not has_tracing_enabled(client.options):
+        return False, None, None, None
+
+    propagation_context = scope.get_active_propagation_context()
+
+    sample_rand = None
+    if propagation_context.baggage is not None:
+        sample_rand = propagation_context.baggage._sample_rand()
+    if sample_rand is None:
+        sample_rand = _generate_sample_rand(propagation_context.trace_id)
+
+    sampling_context = {
+        "name": name,
+        "trace_id": propagation_context.trace_id,
+        "parent_span_id": propagation_context.parent_span_id,
+        "parent_sampled": propagation_context.parent_sampled,
+        "attributes": attributes or {},
+    }
+
+    # If there's a traces_sampler, use that; otherwise use traces_sample_rate
+    traces_sampler_defined = callable(client.options.get("traces_sampler"))
+    if traces_sampler_defined:
+        sample_rate = client.options["traces_sampler"](sampling_context)
+    else:
+        if sampling_context["parent_sampled"] is not None:
+            sample_rate = sampling_context["parent_sampled"]
+        else:
+            sample_rate = client.options["traces_sample_rate"]
+
+    # Validate whether the sample_rate we got is actually valid. Since
+    # traces_sampler is user-provided, it could return anything.
+    if not is_valid_sample_rate(sample_rate, source="Tracing"):
+        logger.warning(f"[Tracing] Discarding {name} because of invalid sample rate.")
+        return False, None, None, None
+
+    sample_rate = float(sample_rate)
+
+    # Adjust sample rate if we're under backpressure
+    if client.monitor:
+        sample_rate /= 2**client.monitor.downsample_factor
+
+    outcome: "Optional[str]" = None
+
+    if not sample_rate:
+        if traces_sampler_defined:
+            reason = "traces_sampler returned 0 or False"
+        else:
+            reason = "traces_sample_rate is set to 0"
+
+        logger.debug(f"[Tracing] Discarding {name} because {reason}")
+        if client.monitor and client.monitor.downsample_factor > 0:
+            outcome = "backpressure"
+        else:
+            outcome = "sample_rate"
+
+        return False, 0.0, None, outcome
+
+    sampled = sample_rand < sample_rate
+
+    if sampled:
+        logger.debug(f"[Tracing] Starting {name}")
+        outcome = None
+    else:
+        logger.debug(
+            f"[Tracing] Discarding {name} because it's not included in the random sample (sampling rate = {sample_rate})"
+        )
+        outcome = "sample_rate"
+
+    return sampled, sample_rate, sample_rand, outcome
+
+
 # Circular imports
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,