ref: Add sampling to span first (#5617)

Add sampling to span first.

- In span first, we make a sampling decision BEFORE a span is actually
created. Based on the sampling decision (and potentially other filtering
mechanisms like `ignore_spans` -- coming soon) we either create a
`StreamedSpan` or a `NoopStreamedSpan`.
- No-op spans will remember why they were unsampled and will emit a
client report accordingly when finished.
- The logic that makes the sampling decision has been ported from
[`_set_initial_sampling_decision`](https://github.com/getsentry/sentry-python/blob/a04ba6e1e5d67a587962f06837952bea4546193f/sentry_sdk/tracing.py#L1145-L1244).
- The effective sampling rate, which might potentially differ from the
incoming baggage (because e.g. backpressure handling is active and
lowers the effective sample rate), is then updated in the propagation
context (ported from
[here](https://github.com/getsentry/sentry-python/blob/a04ba6e1e5d67a587962f06837952bea4546193f/sentry_sdk/scope.py#L1086-L1092)).

Author: sentrivana

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 @@ def start_streamed_span(
         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 @@ def start_streamed_span(
                 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 @@ def start_streamed_span(
                 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 @@ class StreamedSpan:
         "_scope",
         "_previous_span_on_scope",
         "_baggage",
+        "_sample_rand",
+        "_sample_rate",
     )
 
     def __init__(
@@ -238,6 +240,8 @@ def __init__(
         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 @@ def __init__(
         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
@@ -350,7 +356,9 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
             if isinstance(end_timestamp, datetime):
                 self._timestamp = end_timestamp
             else:
-                logger.debug("Failed to set end_timestamp. Using current time instead.")
+                logger.debug(
+                    "[Tracing] Failed to set end_timestamp. Using current time instead."
+                )
 
         if self._timestamp is None:
             try:
@@ -395,7 +403,7 @@ def status(self, status: "Union[SpanStatus, str]") -> None:
 
         if status not in {e.value for e in SpanStatus}:
             logger.debug(
-                f'Unsupported span status {status}. Expected one of: "ok", "error"'
+                f'[Tracing] Unsupported span status {status}. Expected one of: "ok", "error"'
             )
             return
 
@@ -441,13 +449,20 @@ def timestamp(self) -> "Optional[datetime]":
 
 
 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._finished = False
 
         self._start()
 
@@ -471,16 +486,28 @@ def _start(self) -> None:
         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
+        if self._unsampled_reason is not None:
+            client = sentry_sdk.get_client()
+            if client.is_active() and client.transport:
+                logger.debug(
+                    f"[Tracing] Discarding span because sampled=False (reason: {self._unsampled_reason})"
+                )
+                client.transport.record_lost_event(
+                    reason=self._unsampled_reason,
+                    data_category="span",
+                    quantity=1,
+                )
+
+        if self._scope and hasattr(self, "_previous_span_on_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,89 @@ 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)
+
+    # 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:
+        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": dict(attributes) if attributes else {},
+        }
+
+        sample_rate = client.options["traces_sampler"](sampling_context)
+    else:
+        if propagation_context.parent_sampled is not None:
+            sample_rate = propagation_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, "sample_rate"
+
+    sample_rate = float(sample_rate)
+    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}")
+        return False, 0.0, None, "sample_rate"
+
+    # Adjust sample rate if we're under backpressure
+    if client.monitor:
+        sample_rate /= 2**client.monitor.downsample_factor
+
+        if not sample_rate:
+            logger.debug(f"[Tracing] Discarding {name} because backpressure")
+            return False, 0.0, None, "backpressure"
+
+    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,