ref: Allow to start and finish StreamedSpans

Author: sentrivana

sentry_sdk/traces.py

@@ -6,12 +6,20 @@
 """
 
 import uuid
+import warnings
+from datetime import datetime, timedelta, timezone
 from enum import Enum
 from typing import TYPE_CHECKING
 
 import sentry_sdk
 from sentry_sdk.tracing_utils import Baggage
-from sentry_sdk.utils import format_attribute, logger
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    format_attribute,
+    logger,
+    nanosecond_time,
+    should_be_treated_as_error,
+)
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Optional, ParamSpec, TypeVar, Union
@@ -182,8 +190,13 @@ class StreamedSpan:
         "_parent_span_id",
         "_segment",
         "_parent_sampled",
+        "_start_timestamp",
+        "_start_timestamp_monotonic_ns",
+        "_finished",
+        "_timestamp",
         "_status",
         "_scope",
+        "_previous_span_on_scope",
         "_baggage",
     )
 
@@ -216,11 +229,24 @@ def __init__(
         self._parent_sampled = parent_sampled
         self._baggage = baggage
 
+        self._start_timestamp = datetime.now(timezone.utc)
+        self._timestamp: "Optional[datetime]" = None
+        self._finished: bool = False
+
+        try:
+            # profiling depends on this value and requires that
+            # it is measured in nanoseconds
+            self._start_timestamp_monotonic_ns = nanosecond_time()
+        except AttributeError:
+            pass
+
         self._span_id: "Optional[str]" = None
 
         self._status = SpanStatus.OK.value
         self.set_attribute("sentry.span.source", SegmentSource.CUSTOM.value)
 
+        self._start()
+
     def __repr__(self) -> str:
         return (
             f"<{self.__class__.__name__}("
@@ -237,7 +263,77 @@ def __enter__(self) -> "StreamedSpan":
     def __exit__(
         self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]"
     ) -> None:
-        pass
+        if value is not None and should_be_treated_as_error(ty, value):
+            self.status = SpanStatus.ERROR
+
+        self._end()
+
+    def end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
+        """
+        Finish this span and queue it for sending.
+
+        :param end_timestamp: End timestamp to use instead of current time.
+        :type end_timestamp: "Optional[Union[float, datetime]]"
+        """
+        try:
+            if end_timestamp and self._timestamp is None:
+                if isinstance(end_timestamp, float):
+                    end_timestamp = datetime.fromtimestamp(end_timestamp, timezone.utc)
+                self._timestamp = end_timestamp
+        except AttributeError:
+            logger.debug(f"Failed to set end_timestamp: {end_timestamp}")
+
+        self._end()
+
+    def finish(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
+        warnings.warn(
+            "span.finish() is deprecated. Use span.end() instead.",
+            stacklevel=2,
+            category=DeprecationWarning,
+        )
+
+        self.end(end_timestamp)
+
+    def _start(self) -> None:
+        if self._active:
+            old_span = self._scope.span
+            self._scope.span = self
+            self._previous_span_on_scope = old_span
+
+    def _end(self) -> None:
+        if self._finished is True:
+            # This span is already finished, ignore.
+            return
+
+        # Detach from scope
+        if self._active:
+            with capture_internal_exceptions():
+                old_span = self._previous_span_on_scope
+                del self._previous_span_on_scope
+                self._scope.span = old_span
+
+        client = sentry_sdk.get_client()
+        if not client.is_active():
+            return
+
+        # Set attributes from the segment
+        self.set_attribute("sentry.segment.id", self._segment.span_id)
+        self.set_attribute("sentry.segment.name", self._segment.name)
+
+        # Set the end timestamp if not set yet (e.g. via span.end(<timestamp>))
+        if self._timestamp is None:
+            try:
+                elapsed = nanosecond_time() - self._start_timestamp_monotonic_ns
+                self._timestamp = self._start_timestamp + timedelta(
+                    microseconds=elapsed / 1000
+                )
+            except AttributeError:
+                self._timestamp = datetime.now(timezone.utc)
+
+        self._finished = True
+
+        # Finally, queue the span for sending to Sentry
+        self._scope._capture_span(self)
 
     def get_attributes(self) -> "Attributes":
         return self._attributes
@@ -302,10 +398,27 @@ def trace_id(self) -> str:
     def sampled(self) -> "Optional[bool]":
         return True
 
+    @property
+    def start_timestamp(self) -> "Optional[datetime]":
+        return self._start_timestamp
+
+    @property
+    def timestamp(self) -> "Optional[datetime]":
+        return self._timestamp
+
 
 class NoOpStreamedSpan(StreamedSpan):
-    def __init__(self) -> None:
-        pass
+    __slots__ = (
+        "_scope",
+        "_previous_span_on_scope",
+    )
+
+    def __init__(
+        self,
+        scope: "Optional[sentry_sdk.Scope]" = None,
+    ) -> None:
+        self._scope = scope  # type: ignore[assignment]
+        self._start()
 
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}(sampled={self.sampled})>"
@@ -316,7 +429,24 @@ def __enter__(self) -> "NoOpStreamedSpan":
     def __exit__(
         self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]"
     ) -> None:
-        pass
+        self._end()
+
+    def _start(self) -> None:
+        if self._scope is None:
+            return self
+
+        old_span = self._scope.span
+        self._scope.span = self
+        self._previous_span_on_scope = old_span
+
+    def _end(self) -> None:
+        if self._scope is None:
+            return
+
+        with capture_internal_exceptions():
+            old_span = self._previous_span_on_scope
+            del self._previous_span_on_scope
+            self._scope.span = old_span
 
     def get_attributes(self) -> "Attributes":
         return {}