feat(profiling): Enable profiling for ASGI frameworks (#1824)

This enables profiling for ASGI frameworks. When running in ASGI sync views, the
transaction gets started in the main thread then the request is dispatched to a
handler thread. We want to set the handler thread as the active thread id to
ensure that profiles will show it on first render.

Author: Zylphrex

sentry_sdk/client.py

@@ -433,9 +433,7 @@ def capture_event(
 
             if is_transaction:
                 if profile is not None:
-                    envelope.add_profile(
-                        profile.to_json(event_opt, self.options, scope)
-                    )
+                    envelope.add_profile(profile.to_json(event_opt, self.options))
                 envelope.add_transaction(event_opt)
             else:
                 envelope.add_event(event_opt)

sentry_sdk/integrations/asgi.py

@@ -14,6 +14,7 @@
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.integrations._wsgi_common import _filter_headers
 from sentry_sdk.integrations.modules import _get_installed_modules
+from sentry_sdk.profiler import start_profiling
 from sentry_sdk.sessions import auto_session_tracking
 from sentry_sdk.tracing import (
     SOURCE_FOR_STYLE,
@@ -175,7 +176,7 @@ async def _run_app(self, scope, callback):
 
                     with hub.start_transaction(
                         transaction, custom_sampling_context={"asgi_scope": scope}
-                    ):
+                    ), start_profiling(transaction, hub):
                         # XXX: Would be cool to have correct span status, but we
                         # would have to wrap send(). That is a bit hard to do with
                         # the current abstraction over ASGI 2/3.

sentry_sdk/integrations/django/asgi.py

@@ -7,6 +7,7 @@
 """
 
 import asyncio
+import threading
 
 from sentry_sdk import Hub, _functools
 from sentry_sdk._types import MYPY
@@ -89,10 +90,14 @@ def wrap_async_view(hub, callback):
     async def sentry_wrapped_callback(request, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
 
-        with hub.start_span(
-            op=OP.VIEW_RENDER, description=request.resolver_match.view_name
-        ):
-            return await callback(request, *args, **kwargs)
+        with hub.configure_scope() as sentry_scope:
+            if sentry_scope.profile is not None:
+                sentry_scope.profile.active_thread_id = threading.current_thread().ident
+
+            with hub.start_span(
+                op=OP.VIEW_RENDER, description=request.resolver_match.view_name
+            ):
+                return await callback(request, *args, **kwargs)
 
     return sentry_wrapped_callback
 

sentry_sdk/integrations/django/views.py

@@ -1,3 +1,5 @@
+import threading
+
 from sentry_sdk.consts import OP
 from sentry_sdk.hub import Hub
 from sentry_sdk._types import MYPY
@@ -73,9 +75,15 @@ def _wrap_sync_view(hub, callback):
     @_functools.wraps(callback)
     def sentry_wrapped_callback(request, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
-        with hub.start_span(
-            op=OP.VIEW_RENDER, description=request.resolver_match.view_name
-        ):
-            return callback(request, *args, **kwargs)
+        with hub.configure_scope() as sentry_scope:
+            # set the active thread id to the handler thread for sync views
+            # this isn't necessary for async views since that runs on main
+            if sentry_scope.profile is not None:
+                sentry_scope.profile.active_thread_id = threading.current_thread().ident
+
+            with hub.start_span(
+                op=OP.VIEW_RENDER, description=request.resolver_match.view_name
+            ):
+                return callback(request, *args, **kwargs)
 
     return sentry_wrapped_callback

sentry_sdk/integrations/fastapi.py

@@ -1,3 +1,6 @@
+import asyncio
+import threading
+
 from sentry_sdk._types import MYPY
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.integrations import DidNotEnable
@@ -62,6 +65,26 @@ def patch_get_request_handler():
 
     def _sentry_get_request_handler(*args, **kwargs):
         # type: (*Any, **Any) -> Any
+        dependant = kwargs.get("dependant")
+        if (
+            dependant
+            and dependant.call is not None
+            and not asyncio.iscoroutinefunction(dependant.call)
+        ):
+            old_call = dependant.call
+
+            def _sentry_call(*args, **kwargs):
+                # type: (*Any, **Any) -> Any
+                hub = Hub.current
+                with hub.configure_scope() as sentry_scope:
+                    if sentry_scope.profile is not None:
+                        sentry_scope.profile.active_thread_id = (
+                            threading.current_thread().ident
+                        )
+                    return old_call(*args, **kwargs)
+
+            dependant.call = _sentry_call
+
         old_app = old_get_request_handler(*args, **kwargs)
 
         async def _sentry_app(*args, **kwargs):

sentry_sdk/integrations/starlette.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import functools
+import threading
 
 from sentry_sdk._compat import iteritems
 from sentry_sdk._types import MYPY
@@ -403,6 +404,11 @@ def _sentry_sync_func(*args, **kwargs):
                     return old_func(*args, **kwargs)
 
                 with hub.configure_scope() as sentry_scope:
+                    if sentry_scope.profile is not None:
+                        sentry_scope.profile.active_thread_id = (
+                            threading.current_thread().ident
+                        )
+
                     request = args[0]
 
                     _set_transaction_name_and_source(

sentry_sdk/profiler.py

@@ -46,7 +46,6 @@
     from typing import Sequence
     from typing import Tuple
     from typing_extensions import TypedDict
-    import sentry_sdk.scope
     import sentry_sdk.tracing
 
     ThreadId = str
@@ -329,10 +328,13 @@ def __init__(
         self,
         scheduler,  # type: Scheduler
         transaction,  # type: sentry_sdk.tracing.Transaction
+        hub=None,  # type: Optional[sentry_sdk.Hub]
     ):
         # type: (...) -> None
         self.scheduler = scheduler
         self.transaction = transaction
+        self.hub = hub
+        self.active_thread_id = None  # type: Optional[int]
         self.start_ns = 0  # type: int
         self.stop_ns = 0  # type: int
         self.active = False  # type: bool
@@ -347,6 +349,14 @@ def __init__(
 
     def __enter__(self):
         # type: () -> None
+        hub = self.hub or sentry_sdk.Hub.current
+
+        _, scope = hub._stack[-1]
+        old_profile = scope.profile
+        scope.profile = self
+
+        self._context_manager_state = (hub, scope, old_profile)
+
         self.start_ns = nanosecond_time()
         self.scheduler.start_profiling(self)
 
@@ -355,6 +365,11 @@ def __exit__(self, ty, value, tb):
         self.scheduler.stop_profiling(self)
         self.stop_ns = nanosecond_time()
 
+        _, scope, old_profile = self._context_manager_state
+        del self._context_manager_state
+
+        scope.profile = old_profile
+
     def write(self, ts, sample):
         # type: (int, RawSample) -> None
         if ts < self.start_ns:
@@ -414,18 +429,14 @@ def process(self):
             "thread_metadata": thread_metadata,
         }
 
-    def to_json(self, event_opt, options, scope):
-        # type: (Any, Dict[str, Any], Optional[sentry_sdk.scope.Scope]) -> Dict[str, Any]
-
+    def to_json(self, event_opt, options):
+        # type: (Any, Dict[str, Any]) -> Dict[str, Any]
         profile = self.process()
 
         handle_in_app_impl(
             profile["frames"], options["in_app_exclude"], options["in_app_include"]
         )
 
-        # the active thread id from the scope always take priorty if it exists
-        active_thread_id = None if scope is None else scope.active_thread_id
-
         return {
             "environment": event_opt.get("environment"),
             "event_id": uuid.uuid4().hex,
@@ -459,8 +470,8 @@ def to_json(self, event_opt, options, scope):
                     "trace_id": self.transaction.trace_id,
                     "active_thread_id": str(
                         self.transaction._active_thread_id
-                        if active_thread_id is None
-                        else active_thread_id
+                        if self.active_thread_id is None
+                        else self.active_thread_id
                     ),
                 }
             ],
@@ -739,7 +750,7 @@ def start_profiling(transaction, hub=None):
     # if profiling was not enabled, this should be a noop
     if _should_profile(transaction, hub):
         assert _scheduler is not None
-        with Profile(_scheduler, transaction):
+        with Profile(_scheduler, transaction, hub):
             yield
     else:
         yield

sentry_sdk/scope.py

@@ -27,6 +27,7 @@
         Type,
     )
 
+    from sentry_sdk.profiler import Profile
     from sentry_sdk.tracing import Span
     from sentry_sdk.session import Session
 
@@ -94,10 +95,7 @@ class Scope(object):
         "_session",
         "_attachments",
         "_force_auto_session_tracking",
-        # The thread that is handling the bulk of the work. This can just
-        # be the main thread, but that's not always true. For web frameworks,
-        # this would be the thread handling the request.
-        "_active_thread_id",
+        "_profile",
     )
 
     def __init__(self):
@@ -129,7 +127,7 @@ def clear(self):
         self._session = None  # type: Optional[Session]
         self._force_auto_session_tracking = None  # type: Optional[bool]
 
-        self._active_thread_id = None  # type: Optional[int]
+        self._profile = None  # type: Optional[Profile]
 
     @_attr_setter
     def level(self, value):
@@ -235,15 +233,15 @@ def span(self, span):
                 self._transaction = transaction.name
 
     @property
-    def active_thread_id(self):
-        # type: () -> Optional[int]
-        """Get/set the current active thread id."""
-        return self._active_thread_id
+    def profile(self):
+        # type: () -> Optional[Profile]
+        return self._profile
 
-    def set_active_thread_id(self, active_thread_id):
-        # type: (Optional[int]) -> None
-        """Set the current active thread id."""
-        self._active_thread_id = active_thread_id
+    @profile.setter
+    def profile(self, profile):
+        # type: (Optional[Profile]) -> None
+
+        self._profile = profile
 
     def set_tag(
         self,
@@ -464,8 +462,8 @@ def update_from_scope(self, scope):
             self._span = scope._span
         if scope._attachments:
             self._attachments.extend(scope._attachments)
-        if scope._active_thread_id is not None:
-            self._active_thread_id = scope._active_thread_id
+        if scope._profile:
+            self._profile = scope._profile
 
     def update_from_kwargs(
         self,
@@ -515,7 +513,7 @@ def __copy__(self):
         rv._force_auto_session_tracking = self._force_auto_session_tracking
         rv._attachments = list(self._attachments)
 
-        rv._active_thread_id = self._active_thread_id
+        rv._profile = self._profile
 
         return rv
 

tests/integrations/django/asgi/test_asgi.py

@@ -1,3 +1,5 @@
+import json
+
 import django
 import pytest
 from channels.testing import HttpCommunicator
@@ -70,6 +72,41 @@ async def test_async_views(sentry_init, capture_events, application):
     }
 
 
+@pytest.mark.parametrize("application", APPS)
+@pytest.mark.parametrize("endpoint", ["/sync/thread_ids", "/async/thread_ids"])
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    django.VERSION < (3, 1), reason="async views have been introduced in Django 3.1"
+)
+async def test_active_thread_id(sentry_init, capture_envelopes, endpoint, application):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"profiles_sample_rate": 1.0},
+    )
+
+    envelopes = capture_envelopes()
+
+    comm = HttpCommunicator(application, "GET", endpoint)
+    response = await comm.get_response()
+    assert response["status"] == 200, response["body"]
+
+    await comm.wait()
+
+    data = json.loads(response["body"])
+
+    envelopes = [envelope for envelope in envelopes]
+    assert len(envelopes) == 1
+
+    profiles = [item for item in envelopes[0].items if item.type == "profile"]
+    assert len(profiles) == 1
+
+    for profile in profiles:
+        transactions = profile.payload.json["transactions"]
+        assert len(transactions) == 1
+        assert str(data["active"]) == transactions[0]["active_thread_id"]
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(
     django.VERSION < (3, 1), reason="async views have been introduced in Django 3.1"

tests/integrations/django/myapp/urls.py

@@ -58,6 +58,7 @@ def path(path, *args, **kwargs):
         views.csrf_hello_not_exempt,
         name="csrf_hello_not_exempt",
     ),
+    path("sync/thread_ids", views.thread_ids_sync, name="thread_ids_sync"),
 ]
 
 # async views
@@ -67,6 +68,11 @@ def path(path, *args, **kwargs):
 if views.my_async_view is not None:
     urlpatterns.append(path("my_async_view", views.my_async_view, name="my_async_view"))
 
+if views.thread_ids_async is not None:
+    urlpatterns.append(
+        path("async/thread_ids", views.thread_ids_async, name="thread_ids_async")
+    )
+
 # rest framework
 try:
     urlpatterns.append(