feat(profiling): Enable profiling on all transactions

Up to now, we've only been profiling WSGI transactions. This change will enable
profiling for all transactions.

Author: Zylphrex

sentry_sdk/hub.py

@@ -547,6 +547,7 @@ def start_transaction(
         }
         sampling_context.update(custom_sampling_context)
         transaction._set_initial_sampling_decision(sampling_context=sampling_context)
+        transaction._set_profiling_decision()
 
         # we don't bother to keep spans if we already know we're not going to
         # send the transaction

sentry_sdk/integrations/asgi.py

@@ -14,7 +14,6 @@
 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,
@@ -176,7 +175,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/wsgi.py

@@ -12,7 +12,6 @@
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_ROUTE
 from sentry_sdk.sessions import auto_session_tracking
 from sentry_sdk.integrations._wsgi_common import _filter_headers
-from sentry_sdk.profiler import start_profiling
 
 from sentry_sdk._types import MYPY
 
@@ -132,7 +131,7 @@ def __call__(self, environ, start_response):
 
                     with hub.start_transaction(
                         transaction, custom_sampling_context={"wsgi_environ": environ}
-                    ), start_profiling(transaction, hub):
+                    ):
                         try:
                             rv = self.app(
                                 environ,

sentry_sdk/profiler.py

@@ -347,7 +347,7 @@ def __init__(
 
         transaction._profile = self
 
-    def __enter__(self):
+    def start(self):
         # type: () -> None
         hub = self.hub or sentry_sdk.Hub.current
 
@@ -360,8 +360,8 @@ def __enter__(self):
         self.start_ns = nanosecond_time()
         self.scheduler.start_profiling(self)
 
-    def __exit__(self, ty, value, tb):
-        # type: (Optional[Any], Optional[Any], Optional[Any]) -> None
+    def stop(self):
+        # type: () -> None
         self.scheduler.stop_profiling(self)
         self.stop_ns = nanosecond_time()
 
@@ -370,6 +370,14 @@ def __exit__(self, ty, value, tb):
 
         scope.profile = old_profile
 
+    def __enter__(self):
+        # type: () -> None
+        self.start()
+
+    def __exit__(self, ty, value, tb):
+        # type: (Optional[Any], Optional[Any], Optional[Any]) -> None
+        self.stop()
+
     def write(self, ts, sample):
         # type: (int, RawSample) -> None
         if ts < self.start_ns:
@@ -477,6 +485,15 @@ def to_json(self, event_opt, options):
             ],
         }
 
+    @staticmethod
+    def from_transaction(transaction, hub):
+        # type: (sentry_sdk.tracing.Transaction, Optional[sentry_sdk.Hub]) -> Optional[Profile]
+        if not _should_profile(transaction, hub):
+            return None
+
+        assert _scheduler is not None
+        return Profile(_scheduler, transaction, hub)
+
 
 class Scheduler(object):
     mode = "unknown"
@@ -714,7 +731,7 @@ def run(self):
 
 
 def _should_profile(transaction, hub):
-    # type: (sentry_sdk.tracing.Transaction, sentry_sdk.Hub) -> bool
+    # type: (sentry_sdk.tracing.Transaction, Optional[sentry_sdk.Hub]) -> bool
 
     # The corresponding transaction was not sampled,
     # so don't generate a profile for it.
@@ -725,6 +742,7 @@ def _should_profile(transaction, hub):
     if _scheduler is None:
         return False
 
+    hub = hub or sentry_sdk.Hub.current
     client = hub.client
 
     # The client is None, so we can't get the sample rate.
@@ -740,17 +758,3 @@ def _should_profile(transaction, hub):
         return False
 
     return random.random() < float(profiles_sample_rate)
-
-
-@contextmanager
-def start_profiling(transaction, hub=None):
-    # type: (sentry_sdk.tracing.Transaction, Optional[sentry_sdk.Hub]) -> Generator[None, None, None]
-    hub = hub or sentry_sdk.Hub.current
-
-    # if profiling was not enabled, this should be a noop
-    if _should_profile(transaction, hub):
-        assert _scheduler is not None
-        with Profile(_scheduler, transaction, hub):
-            yield
-    else:
-        yield

sentry_sdk/tracing.py

@@ -7,6 +7,7 @@
 
 import sentry_sdk
 from sentry_sdk.consts import INSTRUMENTER
+from sentry_sdk.profiler import Profile
 from sentry_sdk.utils import logger
 from sentry_sdk._types import MYPY
 
@@ -21,7 +22,6 @@
     from typing import Tuple
     from typing import Iterator
 
-    import sentry_sdk.profiler
     from sentry_sdk._types import Event, SamplingContext, MeasurementUnit
 
 BAGGAGE_HEADER_NAME = "baggage"
@@ -604,7 +604,7 @@ def __init__(
         self._third_party_tracestate = third_party_tracestate
         self._measurements = {}  # type: Dict[str, Any]
         self._contexts = {}  # type: Dict[str, Any]
-        self._profile = None  # type: Optional[sentry_sdk.profiler.Profile]
+        self._profile = None  # type: Optional[Profile]
         self._baggage = baggage
         # for profiling, we want to know on which thread a transaction is started
         # to accurately show the active thread in the UI
@@ -708,6 +708,7 @@ def finish(self, hub=None, end_timestamp=None):
         }  # type: Event
 
         if hub.client is not None and self._profile is not None:
+            self._profile.stop()
             event["profile"] = self._profile
 
         if has_custom_measurements_enabled():
@@ -853,6 +854,12 @@ def _set_initial_sampling_decision(self, sampling_context):
                 )
             )
 
+    def _set_profiling_decision(self):
+        # type: () -> None
+        self._profile = Profile.from_transaction(self, self.hub)
+        if self._profile is not None:
+            self._profile.start()
+
 
 class NoOpSpan(Span):
     def __repr__(self):