fix(profiling): Update active thread for asgi (#3669)

Ensure the handling thread is set on the transaction for asgi transactions not just main thread.

---------

Co-authored-by: Anton Pirker <[email protected]>

Author: Zylphrex

sentry_sdk/integrations/django/asgi.py

@@ -172,6 +172,10 @@ def wrap_async_view(callback):
     @functools.wraps(callback)
     async def sentry_wrapped_callback(request, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
+        current_scope = sentry_sdk.get_current_scope()
+        if current_scope.transaction is not None:
+            current_scope.transaction.update_active_thread()
+
         sentry_scope = sentry_sdk.get_isolation_scope()
         if sentry_scope.profile is not None:
             sentry_scope.profile.update_active_thread_id()

sentry_sdk/integrations/django/views.py

@@ -76,6 +76,10 @@ def _wrap_sync_view(callback):
     @functools.wraps(callback)
     def sentry_wrapped_callback(request, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
+        current_scope = sentry_sdk.get_current_scope()
+        if current_scope.transaction is not None:
+            current_scope.transaction.update_active_thread()
+
         sentry_scope = sentry_sdk.get_isolation_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

sentry_sdk/integrations/fastapi.py

@@ -88,9 +88,14 @@ def _sentry_get_request_handler(*args, **kwargs):
             @wraps(old_call)
             def _sentry_call(*args, **kwargs):
                 # type: (*Any, **Any) -> Any
+                current_scope = sentry_sdk.get_current_scope()
+                if current_scope.transaction is not None:
+                    current_scope.transaction.update_active_thread()
+
                 sentry_scope = sentry_sdk.get_isolation_scope()
                 if sentry_scope.profile is not None:
                     sentry_scope.profile.update_active_thread_id()
+
                 return old_call(*args, **kwargs)
 
             dependant.call = _sentry_call

sentry_sdk/integrations/quart.py

@@ -1,6 +1,5 @@
 import asyncio
 import inspect
-import threading
 from functools import wraps
 
 import sentry_sdk
@@ -122,11 +121,13 @@ def decorator(old_func):
                 @ensure_integration_enabled(QuartIntegration, old_func)
                 def _sentry_func(*args, **kwargs):
                     # type: (*Any, **Any) -> Any
-                    scope = sentry_sdk.get_isolation_scope()
-                    if scope.profile is not None:
-                        scope.profile.active_thread_id = (
-                            threading.current_thread().ident
-                        )
+                    current_scope = sentry_sdk.get_current_scope()
+                    if current_scope.transaction is not None:
+                        current_scope.transaction.update_active_thread()
+
+                    sentry_scope = sentry_sdk.get_isolation_scope()
+                    if sentry_scope.profile is not None:
+                        sentry_scope.profile.update_active_thread_id()
 
                     return old_func(*args, **kwargs)
 

sentry_sdk/integrations/starlette.py

@@ -487,8 +487,11 @@ def _sentry_sync_func(*args, **kwargs):
                 if integration is None:
                     return old_func(*args, **kwargs)
 
-                sentry_scope = sentry_sdk.get_isolation_scope()
+                current_scope = sentry_sdk.get_current_scope()
+                if current_scope.transaction is not None:
+                    current_scope.transaction.update_active_thread()
 
+                sentry_scope = sentry_sdk.get_isolation_scope()
                 if sentry_scope.profile is not None:
                     sentry_scope.profile.update_active_thread_id()
 

sentry_sdk/tracing.py

@@ -329,8 +329,7 @@ def __init__(
         self._span_recorder = None  # type: Optional[_SpanRecorder]
         self._local_aggregator = None  # type: Optional[LocalAggregator]
 
-        thread_id, thread_name = get_current_thread_meta()
-        self.set_thread(thread_id, thread_name)
+        self.update_active_thread()
         self.set_profiler_id(get_profiler_id())
 
     # TODO this should really live on the Transaction class rather than the Span
@@ -732,6 +731,11 @@ def get_profile_context(self):
             "profiler_id": profiler_id,
         }
 
+    def update_active_thread(self):
+        # type: () -> None
+        thread_id, thread_name = get_current_thread_meta()
+        self.set_thread(thread_id, thread_name)
+
 
 class Transaction(Span):
     """The Transaction is the root element that holds all the spans

tests/integrations/django/asgi/test_asgi.py

@@ -104,14 +104,16 @@ async def test_async_views(sentry_init, capture_events, application):
 @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):
+async def test_active_thread_id(
+    sentry_init, capture_envelopes, teardown_profiling, endpoint, application
+):
     with mock.patch(
         "sentry_sdk.profiler.transaction_profiler.PROFILE_MINIMUM_SAMPLES", 0
     ):
         sentry_init(
             integrations=[DjangoIntegration()],
             traces_sample_rate=1.0,
-            _experiments={"profiles_sample_rate": 1.0},
+            profiles_sample_rate=1.0,
         )
 
         envelopes = capture_envelopes()
@@ -121,17 +123,26 @@ async def test_active_thread_id(sentry_init, capture_envelopes, endpoint, applic
         await comm.wait()
 
         assert response["status"] == 200, response["body"]
-        assert len(envelopes) == 1
 
-        profiles = [item for item in envelopes[0].items if item.type == "profile"]
-        assert len(profiles) == 1
+    assert len(envelopes) == 1
+
+    profiles = [item for item in envelopes[0].items if item.type == "profile"]
+    assert len(profiles) == 1
+
+    data = json.loads(response["body"])
+
+    for item in profiles:
+        transactions = item.payload.json["transactions"]
+        assert len(transactions) == 1
+        assert str(data["active"]) == transactions[0]["active_thread_id"]
 
-        data = json.loads(response["body"])
+    transactions = [item for item in envelopes[0].items if item.type == "transaction"]
+    assert len(transactions) == 1
 
-        for profile in profiles:
-            transactions = profile.payload.json["transactions"]
-            assert len(transactions) == 1
-            assert str(data["active"]) == transactions[0]["active_thread_id"]
+    for item in transactions:
+        transaction = item.payload.json
+        trace_context = transaction["contexts"]["trace"]
+        assert str(data["active"]) == trace_context["data"]["thread.id"]
 
 
 @pytest.mark.asyncio

tests/integrations/fastapi/test_fastapi.py

@@ -184,7 +184,7 @@ def test_legacy_setup(
 def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, endpoint):
     sentry_init(
         traces_sample_rate=1.0,
-        _experiments={"profiles_sample_rate": 1.0},
+        profiles_sample_rate=1.0,
     )
     app = fastapi_app_factory()
     asgi_app = SentryAsgiMiddleware(app)
@@ -203,11 +203,19 @@ def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, en
     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"]
+    for item in profiles:
+        transactions = item.payload.json["transactions"]
         assert len(transactions) == 1
         assert str(data["active"]) == transactions[0]["active_thread_id"]
 
+    transactions = [item for item in envelopes[0].items if item.type == "transaction"]
+    assert len(transactions) == 1
+
+    for item in transactions:
+        transaction = item.payload.json
+        trace_context = transaction["contexts"]["trace"]
+        assert str(data["active"]) == trace_context["data"]["thread.id"]
+
 
 @pytest.mark.asyncio
 async def test_original_request_not_scrubbed(sentry_init, capture_events):