Make profiling work with potel (#3704)

Author: sl0thentr0py

sentry_sdk/integrations/opentelemetry/potel_span_processor.py

@@ -12,14 +12,22 @@
 from opentelemetry.sdk.trace import Span, ReadableSpan, SpanProcessor
 
 from sentry_sdk import capture_event
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.tracing import DEFAULT_SPAN_ORIGIN
+from sentry_sdk.utils import get_current_thread_meta
+from sentry_sdk.profiler.continuous_profiler import (
+    try_autostart_continuous_profiler,
+    get_profiler_id,
+)
+from sentry_sdk.profiler.transaction_profiler import Profile
 from sentry_sdk.integrations.opentelemetry.utils import (
     is_sentry_span,
     convert_from_otel_timestamp,
     extract_span_attributes,
     extract_span_data,
     extract_transaction_name_source,
     get_trace_context,
+    get_profile_context,
     get_sentry_meta,
     set_sentry_meta,
 )
@@ -54,8 +62,11 @@ def __init__(self):
 
     def on_start(self, span, parent_context=None):
         # type: (Span, Optional[Context]) -> None
-        if not is_sentry_span(span):
-            self._add_root_span(span, get_current_span(parent_context))
+        if is_sentry_span(span):
+            return
+
+        self._add_root_span(span, get_current_span(parent_context))
+        self._start_profile(span)
 
     def on_end(self, span):
         # type: (ReadableSpan) -> None
@@ -94,6 +105,32 @@ def _add_root_span(self, span, parent_span):
             # root span points to itself
             set_sentry_meta(span, "root_span", span)
 
+    def _start_profile(self, span):
+        # type: (Span) -> None
+        try_autostart_continuous_profiler()
+        profiler_id = get_profiler_id()
+        thread_id, thread_name = get_current_thread_meta()
+
+        if profiler_id:
+            span.set_attribute(SPANDATA.PROFILER_ID, profiler_id)
+        if thread_id:
+            span.set_attribute(SPANDATA.THREAD_ID, str(thread_id))
+        if thread_name:
+            span.set_attribute(SPANDATA.THREAD_NAME, thread_name)
+
+        is_root_span = not span.parent or span.parent.is_remote
+        sampled = span.context and span.context.trace_flags.sampled
+
+        if is_root_span and sampled:
+            # profiler uses time.perf_counter_ns() so we cannot use the
+            # unix timestamp that is on span.start_time
+            # setting it to 0 means the profiler will internally measure time on start
+            profile = Profile(sampled, 0)
+            # TODO-neel-potel sampling context??
+            profile._set_initial_sampling_decision(sampling_context={})
+            profile.__enter__()
+            set_sentry_meta(span, "profile", profile)
+
     def _flush_root_span(self, span):
         # type: (ReadableSpan) -> None
         transaction_event = self._root_span_to_transaction_event(span)
@@ -147,6 +184,10 @@ def _root_span_to_transaction_event(self, span):
         trace_context = get_trace_context(span, span_data=span_data)
         contexts = {"trace": trace_context}
 
+        profile_context = get_profile_context(span)
+        if profile_context:
+            contexts["profile"] = profile_context
+
         if http_status:
             contexts["response"] = {"status_code": http_status}
 
@@ -162,6 +203,13 @@ def _root_span_to_transaction_event(self, span):
             }
         )
 
+        profile = cast("Optional[Profile]", get_sentry_meta(span, "profile"))
+        if profile:
+            profile.__exit__(None, None, None)
+            if profile.valid():
+                event["profile"] = profile
+                set_sentry_meta(span, "profile", None)
+
         return event
 
     def _span_to_json(self, span):

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -128,6 +128,7 @@ def should_sample(
         has_traces_sampler = callable(client.options.get("traces_sampler"))
         if has_traces_sampler:
             # TODO-anton: Make proper sampling_context
+            # TODO-neel-potel: Make proper sampling_context
             sampling_context = {
                 "transaction_context": {
                     "name": name,

sentry_sdk/integrations/opentelemetry/utils.py

@@ -17,7 +17,7 @@
 
 import sentry_sdk
 from sentry_sdk.utils import Dsn
-from sentry_sdk.consts import SPANSTATUS, OP
+from sentry_sdk.consts import SPANSTATUS, OP, SPANDATA
 from sentry_sdk.tracing import get_span_status_from_http_code, DEFAULT_SPAN_ORIGIN
 from sentry_sdk.tracing_utils import Baggage, LOW_QUALITY_TRANSACTION_SOURCES
 from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
@@ -432,3 +432,15 @@ def set_sentry_meta(span, key, value):
     sentry_meta = getattr(span, "_sentry_meta", {})
     sentry_meta[key] = value
     span._sentry_meta = sentry_meta
+
+
+def get_profile_context(span):
+    # type: (ReadableSpan) -> Optional[dict[str, str]]
+    if not span.attributes:
+        return None
+
+    profiler_id = cast("Optional[str]", span.attributes.get(SPANDATA.PROFILER_ID))
+    if profiler_id is None:
+        return None
+
+    return {"profiler_id": profiler_id}

sentry_sdk/tracing.py

@@ -1549,10 +1549,10 @@ def set_thread(self, thread_id, thread_name):
             if thread_name is not None:
                 self.set_data(SPANDATA.THREAD_NAME, thread_name)
 
-    def set_profiler_id(self, profiler_id):
-        # type: (Optional[str]) -> None
-        if profiler_id is not None:
-            self.set_data(SPANDATA.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)
 
     def set_http_status(self, http_status):
         # type: (int) -> None
@@ -1576,6 +1576,7 @@ def finish(self, end_timestamp=None):
 
     def to_json(self):
         # type: () -> dict[str, Any]
+        # TODO-neel-potel for sampling context
         pass
 
     def get_trace_context(self):
@@ -1589,10 +1590,6 @@ def get_trace_context(self):
 
         return get_trace_context(self._otel_span)
 
-    def get_profile_context(self):
-        # type: () -> Optional[ProfileContext]
-        pass
-
     def set_context(self, key, value):
         # type: (str, Any) -> None
         from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute