feat: Span streaming & new span API

Author: sentrivana

sentry_sdk/_span_batcher.py

@@ -0,0 +1,136 @@
+import threading
+from collections import defaultdict
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from sentry_sdk._batcher import Batcher
+from sentry_sdk.consts import SPANSTATUS
+from sentry_sdk.envelope import Envelope, Item, PayloadRef
+from sentry_sdk.utils import serialize_attribute, safe_repr
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional
+    from sentry_sdk.trace import Span
+
+
+class SpanBatcher(Batcher["Span"]):
+    # TODO[span-first]: size-based flushes
+    MAX_BEFORE_FLUSH = 1000
+    MAX_BEFORE_DROP = 5000
+    FLUSH_WAIT_TIME = 5.0
+
+    TYPE = "span"
+    CONTENT_TYPE = "application/vnd.sentry.items.span.v2+json"
+
+    def __init__(
+        self,
+        capture_func: "Callable[[Envelope], None]",
+        record_lost_func: "Callable[..., None]",
+    ) -> None:
+        # Spans from different traces cannot be emitted in the same envelope
+        # since the envelope contains a shared trace header. That's why we bucket
+        # by trace_id, so that we can then send the buckets each in its own
+        # envelope.
+        # trace_id -> span buffer
+        self._span_buffer: dict[str, list["Span"]] = defaultdict(list)
+        self._capture_func = capture_func
+        self._record_lost_func = record_lost_func
+        self._running = True
+        self._lock = threading.Lock()
+
+        self._flush_event: "threading.Event" = threading.Event()
+
+        self._flusher: "Optional[threading.Thread]" = None
+        self._flusher_pid: "Optional[int]" = None
+
+    def get_size(self) -> int:
+        # caller is responsible for locking before checking this
+        return sum(len(buffer) for buffer in self._span_buffer.values())
+
+    def add(self, span: Span) -> None:
+        if not self._ensure_thread() or self._flusher is None:
+            return None
+
+        with self._lock:
+            size = self.get_size()
+            if size >= self.MAX_BEFORE_DROP:
+                self._record_lost_func(
+                    reason="queue_overflow",
+                    data_category="span",
+                    quantity=1,
+                )
+                return None
+
+            self._span_buffer[span.trace_id].append(span)
+            if size + 1 >= self.MAX_BEFORE_FLUSH:
+                self._flush_event.set()
+
+    @staticmethod
+    def _to_transport_format(item: "Span") -> "Any":
+        is_segment = item.containing_transaction == item
+
+        res = {
+            "trace_id": item.trace_id,
+            "span_id": item.span_id,
+            "name": item.name if is_segment else item.description,
+            "status": SPANSTATUS.OK
+            if item.status == SPANSTATUS.OK
+            else SPANSTATUS.INTERNAL_ERROR,
+            "is_segment": is_segment,
+            "start_timestamp": item.start_timestamp.timestamp(),  # TODO[span-first]
+            "end_timestamp": item.timestamp.timestamp(),
+        }
+
+        if item.parent_span_id:
+            res["parent_span_id"] = item.parent_span_id
+
+        if item._attributes:
+            res["attributes"] = {
+                k: serialize_attribute(v) for (k, v) in item._attributes.items()
+            }
+
+        return res
+
+    def _flush(self):
+        # type: (...) -> Optional[Envelope]
+        from sentry_sdk.utils import format_timestamp
+
+        with self._lock:
+            if len(self._span_buffer) == 0:
+                return None
+
+            for trace_id, spans in self._span_buffer.items():
+                if spans:
+                    trace_context = spans[0].get_trace_context()
+                    dsc = trace_context.get("dynamic_sampling_context")
+                    # XXX[span-first]: empty dsc?
+
+                    envelope = Envelope(
+                        headers={
+                            "sent_at": format_timestamp(datetime.now(timezone.utc)),
+                            "trace": dsc,
+                        }
+                    )
+
+                    envelope.add_item(
+                        Item(
+                            type="span",
+                            content_type="application/vnd.sentry.items.span.v2+json",
+                            headers={
+                                "item_count": len(spans),
+                            },
+                            payload=PayloadRef(
+                                json={
+                                    "items": [
+                                        self._to_transport_format(span)
+                                        for span in spans
+                                    ]
+                                }
+                            ),
+                        )
+                    )
+
+            self._span_buffer.clear()
+
+        self._capture_func(envelope)
+        return envelope

sentry_sdk/_types.py

@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, TypeVar, Union
 
+from sentry_sdk.consts import SPANSTATUS
+
 
 # Re-exported for compat, since code out there in the wild might use this variable.
 MYPY = TYPE_CHECKING
@@ -274,6 +276,26 @@ class SDKInfo(TypedDict):
 
     MetricProcessor = Callable[[Metric, Hint], Optional[Metric]]
 
+    SpanV2Status = Literal[SPANSTATUS.OK, SPANSTATUS.ERROR]
+    # This is the V2 span format
+    # https://develop.sentry.dev/sdk/telemetry/spans/span-protocol/
+    SpanV2 = TypedDict(
+        "SpanV2",
+        {
+            "trace_id": str,
+            "span_id": str,
+            "parent_span_id": Optional[str],
+            "name": str,
+            "status": SpanV2Status,
+            "is_segment": bool,
+            "start_timestamp": float,
+            "end_timestamp": float,
+            "attributes": Attributes,
+        },
+    )
+
+    TraceLifecycleMode = Literal["static", "stream"]
+
     # TODO: Make a proper type definition for this (PRs welcome!)
     Breadcrumb = Dict[str, Any]
 

sentry_sdk/client.py

@@ -11,6 +11,7 @@
 import sentry_sdk
 from sentry_sdk._compat import PY37, check_uwsgi_thread_support
 from sentry_sdk._metrics_batcher import MetricsBatcher
+from sentry_sdk._span_batcher import SpanBatcher
 from sentry_sdk.utils import (
     AnnotatedValue,
     ContextVar,
@@ -31,6 +32,7 @@
 )
 from sentry_sdk.serializer import serialize
 from sentry_sdk.tracing import trace
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.transport import BaseHttpTransport, make_transport
 from sentry_sdk.consts import (
     SPANDATA,
@@ -67,6 +69,7 @@
     from sentry_sdk.scope import Scope
     from sentry_sdk.session import Session
     from sentry_sdk.spotlight import SpotlightClient
+    from sentry_sdk.trace import StreamedSpan
     from sentry_sdk.transport import Transport, Item
     from sentry_sdk._log_batcher import LogBatcher
     from sentry_sdk._metrics_batcher import MetricsBatcher
@@ -188,6 +191,7 @@ def __init__(self, options: "Optional[Dict[str, Any]]" = None) -> None:
         self.monitor: "Optional[Monitor]" = None
         self.log_batcher: "Optional[LogBatcher]" = None
         self.metrics_batcher: "Optional[MetricsBatcher]" = None
+        self.span_batcher: "Optional[SpanBatcher]" = None
         self.integrations: "dict[str, Integration]" = {}
 
     def __getstate__(self, *args: "Any", **kwargs: "Any") -> "Any":
@@ -399,6 +403,12 @@ def _record_lost_event(
                     record_lost_func=_record_lost_event,
                 )
 
+            self.span_batcher = None
+            if has_span_streaming_enabled(self.options):
+                self.span_batcher = SpanBatcher(
+                    capture_func=_capture_envelope, record_lost_func=_record_lost_event
+                )
+
             max_request_body_size = ("always", "never", "small", "medium")
             if self.options["max_request_body_size"] not in max_request_body_size:
                 raise ValueError(
@@ -909,7 +919,10 @@ def capture_event(
         return return_value
 
     def _capture_telemetry(
-        self, telemetry: "Optional[Union[Log, Metric]]", ty: str, scope: "Scope"
+        self,
+        telemetry: "Optional[Union[Log, Metric, StreamedSpan]]",
+        ty: str,
+        scope: "Scope",
     ) -> None:
         # Capture attributes-based telemetry (logs, metrics, spansV2)
         if telemetry is None:
@@ -921,7 +934,7 @@ def _capture_telemetry(
         if ty == "log":
             before_send = get_before_send_log(self.options)
         elif ty == "metric":
-            before_send = get_before_send_metric(self.options)  # type: ignore
+            before_send = get_before_send_metric(self.options)
 
         if before_send is not None:
             telemetry = before_send(telemetry, {})  # type: ignore
@@ -933,7 +946,9 @@ def _capture_telemetry(
         if ty == "log":
             batcher = self.log_batcher
         elif ty == "metric":
-            batcher = self.metrics_batcher  # type: ignore
+            batcher = self.metrics_batcher
+        elif ty == "span":
+            batcher = self.span_batcher
 
         if batcher is not None:
             batcher.add(telemetry)  # type: ignore
@@ -944,6 +959,9 @@ def _capture_log(self, log: "Optional[Log]", scope: "Scope") -> None:
     def _capture_metric(self, metric: "Optional[Metric]", scope: "Scope") -> None:
         self._capture_telemetry(metric, "metric", scope)
 
+    def _capture_span(self, span: "Optional[StreamedSpan]", scope: "Scope") -> None:
+        self._capture_telemetry(span, "span", scope)
+
     def capture_session(
         self,
         session: "Session",

sentry_sdk/consts.py

@@ -2,6 +2,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 
+from sentry_sdk._types import TraceLifecycleMode
+
 # up top to prevent circular import due to integration import
 # This is more or less an arbitrary large-ish value for now, so that we allow
 # pretty long strings (like LLM prompts), but still have *some* upper limit
@@ -82,6 +84,7 @@ class CompressionAlgo(Enum):
             "before_send_log": Optional[Callable[[Log, Hint], Optional[Log]]],
             "enable_metrics": Optional[bool],
             "before_send_metric": Optional[Callable[[Metric, Hint], Optional[Metric]]],
+            "trace_lifecycle": Optional[TraceLifecycleMode],
         },
         total=False,
     )
@@ -877,6 +880,8 @@ class SPANSTATUS:
     UNIMPLEMENTED = "unimplemented"
     UNKNOWN_ERROR = "unknown_error"
 
+    ERROR = "error"  # span-first specific
+
 
 class OP:
     ANTHROPIC_MESSAGES_CREATE = "ai.messages.create.anthropic"

sentry_sdk/envelope.py

@@ -253,6 +253,8 @@ def data_category(self) -> "EventDataCategory":
             return "session"
         elif ty == "attachment":
             return "attachment"
+        elif ty == "span":
+            return "span"
         elif ty == "transaction":
             return "transaction"
         elif ty == "event":

sentry_sdk/scope.py

@@ -32,6 +32,7 @@
     normalize_incoming_data,
     PropagationContext,
 )
+from sentry_sdk.trace import StreamedSpan
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,
     SENTRY_TRACE_HEADER_NAME,
@@ -1147,6 +1148,40 @@ def start_span(
 
             return span
 
+    def start_streamed_span(
+        self,
+        name: str,
+        attributes: "Optional[Attributes]" = None,
+        parent_span: "Optional[StreamedSpan]" = None,
+    ) -> "StreamedSpan":
+        # TODO: rename to start_span once we drop the old API
+        with new_scope():
+            if parent_span is None:
+                # get current span or transaction
+                parent_span = self.span or self.get_isolation_scope().span
+
+            if parent_span is None:
+                # New spans get the `trace_id` from the scope
+                propagation_context = self.get_active_propagation_context()
+                span = StreamedSpan(
+                    name=name,
+                    attributes=attributes,
+                    trace_id=propagation_context.trace_id,
+                    scope=self,
+                )
+            else:
+                # Children take propagation context from the parent span
+                span = StreamedSpan(
+                    name=name,
+                    attributes=attributes,
+                    trace_id=parent_span.trace_id,
+                    parent_span_id=parent_span.span_id,
+                    segment=parent_span.segment,
+                    scope=self,
+                )
+
+            return span
+
     def continue_trace(
         self,
         environ_or_headers: "Dict[str, Any]",
@@ -1180,6 +1215,9 @@ def continue_trace(
             **optional_kwargs,
         )
 
+    def set_propagation_context(self, environ_or_headers: "dict[str, Any]") -> None:
+        self.generate_propagation_context(environ_or_headers)
+
     def capture_event(
         self,
         event: "Event",