From b1290c60208997b082287c724454949ae0166b54 Mon Sep 17 00:00:00 2001
From: Tony Xiao <txiao@sentry.io>
Date: Wed, 7 Dec 2022 06:11:24 -0800
Subject: [PATCH] feat(profiling): Introduce active thread id on scope (#1764)

Up to this point, simply taking the current thread when the transaction/profile
was started was good enough. When using ASGI apps with non async handlers, the
request is received on the main thread. This is also where the transaction or
profile was started. However, the request is handled on another thread using a
thread pool. To support this use case, we want to be able to set the active
thread id on the scope where we can read it when we need it to allow the active
thread id to be set elsewhere.
---
 sentry_sdk/client.py   |  4 +++-
 sentry_sdk/profiler.py | 14 +++++++++++---
 sentry_sdk/scope.py    | 21 +++++++++++++++++++++
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py
index 8af7003156..d32d014d96 100644
--- a/sentry_sdk/client.py
+++ b/sentry_sdk/client.py
@@ -433,7 +433,9 @@ def capture_event(
 
             if is_transaction:
                 if profile is not None:
-                    envelope.add_profile(profile.to_json(event_opt, self.options))
+                    envelope.add_profile(
+                        profile.to_json(event_opt, self.options, scope)
+                    )
                 envelope.add_transaction(event_opt)
             else:
                 envelope.add_event(event_opt)
diff --git a/sentry_sdk/profiler.py b/sentry_sdk/profiler.py
index b38b7af962..21313c9f73 100644
--- a/sentry_sdk/profiler.py
+++ b/sentry_sdk/profiler.py
@@ -51,6 +51,7 @@
     from typing import Sequence
     from typing import Tuple
     from typing_extensions import TypedDict
+    import sentry_sdk.scope
     import sentry_sdk.tracing
 
     RawStack = Tuple[RawFrameData, ...]
@@ -267,8 +268,8 @@ def __exit__(self, ty, value, tb):
         self.scheduler.stop_profiling()
         self._stop_ns = nanosecond_time()
 
-    def to_json(self, event_opt, options):
-        # type: (Any, Dict[str, Any]) -> Dict[str, Any]
+    def to_json(self, event_opt, options, scope):
+        # type: (Any, Dict[str, Any], Optional[sentry_sdk.scope.Scope]) -> Dict[str, Any]
         assert self._start_ns is not None
         assert self._stop_ns is not None
 
@@ -280,6 +281,9 @@ def to_json(self, event_opt, options):
             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,
@@ -311,7 +315,11 @@ def to_json(self, event_opt, options):
                     # because we end the transaction after the profile
                     "relative_end_ns": str(self._stop_ns - self._start_ns),
                     "trace_id": self.transaction.trace_id,
-                    "active_thread_id": str(self.transaction._active_thread_id),
+                    "active_thread_id": str(
+                        self.transaction._active_thread_id
+                        if active_thread_id is None
+                        else active_thread_id
+                    ),
                 }
             ],
         }
diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
index e0a2dc7a8d..f5ac270914 100644
--- a/sentry_sdk/scope.py
+++ b/sentry_sdk/scope.py
@@ -94,6 +94,10 @@ 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",
     )
 
     def __init__(self):
@@ -125,6 +129,8 @@ 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]
+
     @_attr_setter
     def level(self, value):
         # type: (Optional[str]) -> None
@@ -228,6 +234,17 @@ def span(self, span):
             if transaction.name:
                 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 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
+
     def set_tag(
         self,
         key,  # type: str
@@ -447,6 +464,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
 
     def update_from_kwargs(
         self,
@@ -496,6 +515,8 @@ 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
+
         return rv
 
     def __repr__(self):