Fix spans for streaming responses in WSGI based frameworks

sentry_sdk/integrations/wsgi.py

@@ -1,19 +1,20 @@
+from datetime import datetime, timezone
 import sys
 from functools import partial
+from threading import Timer
 
 import sentry_sdk
 from sentry_sdk._werkzeug import get_host, _get_headers
 from sentry_sdk.api import continue_trace
 from sentry_sdk.consts import OP
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import should_send_default_pii, use_isolation_scope, use_scope
 from sentry_sdk.integrations._wsgi_common import (
     DEFAULT_HTTP_METHODS_TO_CAPTURE,
     _filter_headers,
-    nullcontext,
 )
 from sentry_sdk.sessions import track_session
-from sentry_sdk.scope import use_isolation_scope
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_ROUTE
+from sentry_sdk.tracing_utils import finish_running_transaction
 from sentry_sdk.utils import (
     ContextVar,
     capture_internal_exceptions,
@@ -46,6 +47,9 @@ def __call__(self, status, response_headers, exc_info=None):  # type: ignore
             pass
 
 
+MAX_TRANSACTION_DURATION_SECONDS = 5 * 60
+
+
 _wsgi_middleware_applied = ContextVar("sentry_wsgi_middleware_applied")
 
 
@@ -98,6 +102,7 @@ def __call__(self, environ, start_response):
         _wsgi_middleware_applied.set(True)
         try:
             with sentry_sdk.isolation_scope() as scope:
+                current_scope = sentry_sdk.get_current_scope()
                 with track_session(scope, session_mode="request"):
                     with capture_internal_exceptions():
                         scope.clear_breadcrumbs()
@@ -109,6 +114,7 @@ def __call__(self, environ, start_response):
                         )
 
                     method = environ.get("REQUEST_METHOD", "").upper()
+
                     transaction = None
                     if method in self.http_methods_to_capture:
                         transaction = continue_trace(
@@ -119,27 +125,43 @@ def __call__(self, environ, start_response):
                             origin=self.span_origin,
                         )
 
-                    with (
+                    timer = None
+                    if transaction is not None:
                         sentry_sdk.start_transaction(
                             transaction,
                             custom_sampling_context={"wsgi_environ": environ},
+                        ).__enter__()
+                        timer = Timer(
+                            MAX_TRANSACTION_DURATION_SECONDS,
+                            _finish_long_running_transaction,
+                            args=(current_scope, scope),
                         )
-                        if transaction is not None
-                        else nullcontext()
-                    ):
-                        try:
-                            response = self.app(
-                                environ,
-                                partial(
-                                    _sentry_start_response, start_response, transaction
-                                ),
-                            )
-                        except BaseException:
-                            reraise(*_capture_exception())
+                        timer.start()
+
+                    try:
+                        response = self.app(
+                            environ,
+                            partial(
+                                _sentry_start_response,
+                                start_response,
+                                transaction,
+                            ),
+                        )
+                    except BaseException:
+                        exc_info = sys.exc_info()
+                        _capture_exception(exc_info)
+                        finish_running_transaction(current_scope, exc_info, timer)
+                        reraise(*exc_info)
+
         finally:
             _wsgi_middleware_applied.set(False)
 
-        return _ScopedResponse(scope, response)
+        return _ScopedResponse(
+            response=response,
+            current_scope=current_scope,
+            isolation_scope=scope,
+            timer=timer,
+        )
 
 
 def _sentry_start_response(  # type: ignore
@@ -201,13 +223,13 @@ def get_client_ip(environ):
     return environ.get("REMOTE_ADDR")
 
 
-def _capture_exception():
-    # type: () -> ExcInfo
+def _capture_exception(exc_info=None):
+    # type: (Optional[ExcInfo]) -> ExcInfo
     """
     Captures the current exception and sends it to Sentry.
     Returns the ExcInfo tuple to it can be reraised afterwards.
     """
-    exc_info = sys.exc_info()
+    exc_info = exc_info or sys.exc_info()
     e = exc_info[1]
 
     # SystemExit(0) is the only uncaught exception that is expected behavior
@@ -225,7 +247,7 @@ def _capture_exception():
 
 class _ScopedResponse:
     """
-    Users a separate scope for each response chunk.
+    Use separate scopes for each response chunk.
 
     This will make WSGI apps more tolerant against:
     - WSGI servers streaming responses from a different thread/from
@@ -234,37 +256,54 @@ class _ScopedResponse:
     - WSGI servers streaming responses interleaved from the same thread
     """
 
-    __slots__ = ("_response", "_scope")
+    __slots__ = ("_response", "_current_scope", "_isolation_scope", "_timer")
 
-    def __init__(self, scope, response):
-        # type: (sentry_sdk.scope.Scope, Iterator[bytes]) -> None
-        self._scope = scope
+    def __init__(
+        self,
+        response,  # type: Iterator[bytes]
+        current_scope,  # type: sentry_sdk.scope.Scope
+        isolation_scope,  # type: sentry_sdk.scope.Scope
+        timer=None,  # type: Optional[Timer]
+    ):
+        # type: (...) -> None
         self._response = response
+        self._current_scope = current_scope
+        self._isolation_scope = isolation_scope
+        self._timer = timer
 
     def __iter__(self):
         # type: () -> Iterator[bytes]
         iterator = iter(self._response)
 
-        while True:
-            with use_isolation_scope(self._scope):
-                try:
-                    chunk = next(iterator)
-                except StopIteration:
-                    break
-                except BaseException:
-                    reraise(*_capture_exception())
+        try:
+            while True:
+                with use_isolation_scope(self._isolation_scope):
+                    with use_scope(self._current_scope):
+                        try:
+                            chunk = next(iterator)
+                        except StopIteration:
+                            break
+                        except BaseException:
+                            reraise(*_capture_exception())
+
+                yield chunk
 
-            yield chunk
+        finally:
+            with use_isolation_scope(self._isolation_scope):
+                with use_scope(self._current_scope):
+                    finish_running_transaction(timer=self._timer)
 
     def close(self):
         # type: () -> None
-        with use_isolation_scope(self._scope):
-            try:
-                self._response.close()  # type: ignore
-            except AttributeError:
-                pass
-            except BaseException:
-                reraise(*_capture_exception())
+        with use_isolation_scope(self._isolation_scope):
+            with use_scope(self._current_scope):
+                try:
+                    finish_running_transaction(timer=self._timer)
+                    self._response.close()  # type: ignore
+                except AttributeError:
+                    pass
+                except BaseException:
+                    reraise(*_capture_exception())
 
 
 def _make_wsgi_event_processor(environ, use_x_forwarded_for):
@@ -308,3 +347,22 @@ def event_processor(event, hint):
         return event
 
     return event_processor
+
+
+def _finish_long_running_transaction(current_scope, isolation_scope):
+    # type: (sentry_sdk.scope.Scope, sentry_sdk.scope.Scope) -> None
+    """
+    Make sure we don't keep transactions open for too long.
+    Triggered after MAX_TRANSACTION_DURATION_SECONDS have passed.
+    """
+    try:
+        transaction_duration = (
+            datetime.now(timezone.utc) - current_scope.transaction.start_timestamp
+        ).total_seconds()
+        if transaction_duration > MAX_TRANSACTION_DURATION_SECONDS:
+            with use_isolation_scope(isolation_scope):
+                with use_scope(current_scope):
+                    finish_running_transaction()
+    except AttributeError:
+        # transaction is not there anymore
+        pass

sentry_sdk/tracing_utils.py

@@ -36,6 +36,9 @@
 
     from types import FrameType
 
+    from sentry_sdk._types import ExcInfo
+    from threading import Timer
+
 
 SENTRY_TRACE_REGEX = re.compile(
     "^[ \t]*"  # whitespace
@@ -739,3 +742,18 @@ def get_current_span(scope=None):
 
 if TYPE_CHECKING:
     from sentry_sdk.tracing import Span
+
+
+def finish_running_transaction(scope=None, exc_info=None, timer=None):
+    # type: (Optional[sentry_sdk.Scope], Optional[ExcInfo], Optional[Timer]) -> None
+    current_scope = scope or sentry_sdk.get_current_scope()
+    if current_scope.transaction is not None and hasattr(
+        current_scope.transaction, "_context_manager_state"
+    ):
+        if timer is not None:
+            timer.cancel()
+
+        if exc_info is not None:
+            current_scope.transaction.__exit__(*exc_info)
+        else:
+            current_scope.transaction.__exit__(None, None, None)