Fix spans for streaming responses in WSGI based frameworks (#3798)

Fixes spans in streaming responses when using WSGI based frameworks. Only close the transaction once the response was consumed. This way all the spans created during creation of the response will be recorded with the transaction:

- The transaction stays open until all the streaming blocks are sent to the client. (because of this I had to update the tests, to make sure the tests, consume the response, because the Werkzeug test client (used by Flask and Django and our Strawberry tests) will not close the WSGI response)
- A maximum runtime of 5 minutes for transactions is enforced. (like Javascript does it)
- When using a generator to generate the streaming response, it uses the correct scopes to have correct parent-child relationship of spans created in the generator.

People having Sentry in a streaming application will:
- See an increase in their transaction duration to up to 5 minutes
- Get the correct span tree for streaming responses generated by a generator

Fixes #3736

Author: antonpirker

sentry_sdk/integrations/wsgi.py

@@ -1,19 +1,19 @@
 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 +46,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 +101,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 +113,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 +124,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 +222,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 +246,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 +255,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 +346,18 @@ 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:
+        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
+    if timer is not None:
+        timer.cancel()
+
+    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 exc_info is not None:
+            current_scope.transaction.__exit__(*exc_info)
+        else:
+            current_scope.transaction.__exit__(None, None, None)