ref: Add cont profiling support to span first (19)

sentry_sdk/traces.py

@@ -15,6 +15,11 @@
 
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.profiler.continuous_profiler import (
+    get_profiler_id,
+    try_autostart_continuous_profiler,
+    try_profile_lifecycle_trace_start,
+)
 from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -28,6 +33,7 @@
 if TYPE_CHECKING:
     from typing import Any, Callable, Iterator, Optional, ParamSpec, TypeVar, Union
     from sentry_sdk._types import Attributes, AttributeValue
+    from sentry_sdk.profiler.continuous_profiler import ContinuousProfile
 
     P = ParamSpec("P")
     R = TypeVar("R")
@@ -232,6 +238,7 @@ class StreamedSpan:
         "_baggage",
         "_sample_rand",
         "_sample_rate",
+        "_continuous_profile",
     )
 
     def __init__(
@@ -284,6 +291,10 @@ def __init__(
 
         self._update_active_thread()
 
+        self._continuous_profile: "Optional[ContinuousProfile]" = None
+        self._start_profile()
+        self._set_profile_id(get_profiler_id())
+
         self._start()
 
     def __repr__(self) -> str:
@@ -340,6 +351,11 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
             # This span is already finished, ignore.
             return
 
+        # Stop the profiler
+        if self._is_segment() and self._continuous_profile is not None:
+            with capture_internal_exceptions():
+                self._continuous_profile.stop()
+
         # Detach from scope
         if self._active:
             with capture_internal_exceptions():
@@ -528,6 +544,18 @@ def _get_trace_context(self) -> "dict[str, Any]":
 
         return context
 
+    def _set_profile_id(self, profiler_id: "Optional[str]") -> None:
+        if profiler_id is not None:
+            self.set_attribute("sentry.profiler_id", profiler_id)
+
+    def _start_profile(self) -> None:
+        if not self._is_segment():
+            return
+
+        try_autostart_continuous_profiler()
+
+        self._continuous_profile = try_profile_lifecycle_trace_start()
+
 
 class NoOpStreamedSpan(StreamedSpan):
     __slots__ = (

tests/tracing/test_span_streaming.py

@@ -1,12 +1,14 @@
 import asyncio
 import re
 import sys
+import time
 from typing import Any
 from unittest import mock
 
 import pytest
 
 import sentry_sdk
+from sentry_sdk.profiler.continuous_profiler import get_profiler_id
 from sentry_sdk.traces import NoOpStreamedSpan, SpanStatus, StreamedSpan
 
 minimum_python_38 = pytest.mark.skipif(
@@ -1322,6 +1324,102 @@ def test_ignore_spans_reparenting(sentry_init, capture_envelopes):
     assert span5["parent_span_id"] == span3["span_id"]
 
 
+@mock.patch("sentry_sdk.profiler.continuous_profiler.DEFAULT_SAMPLING_FREQUENCY", 21)
+def test_segment_span_has_profiler_id(
+    sentry_init, capture_envelopes, teardown_profiling
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        profile_lifecycle="trace",
+        profiler_mode="thread",
+        profile_session_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+            "continuous_profiling_auto_start": True,
+        },
+    )
+    envelopes = capture_envelopes()
+
+    with sentry_sdk.traces.start_span(name="profiled segment"):
+        time.sleep(0.1)
+
+    sentry_sdk.get_client().flush()
+    time.sleep(0.3)  # wait for profiler to flush
+
+    spans = envelopes_to_spans(envelopes)
+    assert len(spans) == 1
+    assert "sentry.profiler_id" in spans[0]["attributes"]
+
+    profile_chunks = [
+        item
+        for envelope in envelopes
+        for item in envelope.items
+        if item.type == "profile_chunk"
+    ]
+    assert len(profile_chunks) > 0
+
+
+def test_segment_span_no_profiler_id_when_unsampled(
+    sentry_init, capture_envelopes, teardown_profiling
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        profile_lifecycle="trace",
+        profiler_mode="thread",
+        profile_session_sample_rate=0.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+            "continuous_profiling_auto_start": True,
+        },
+    )
+    envelopes = capture_envelopes()
+
+    with sentry_sdk.traces.start_span(name="segment"):
+        time.sleep(0.05)
+
+    sentry_sdk.get_client().flush()
+    time.sleep(0.2)
+
+    spans = envelopes_to_spans(envelopes)
+    assert len(spans) == 1
+    assert "sentry.profiler_id" not in spans[0]["attributes"]
+
+    profile_chunks = [
+        item
+        for envelope in envelopes
+        for item in envelope.items
+        if item.type == "profile_chunk"
+    ]
+    assert len(profile_chunks) == 0
+
+
+@mock.patch("sentry_sdk.profiler.continuous_profiler.DEFAULT_SAMPLING_FREQUENCY", 21)
+def test_profile_stops_when_segment_ends(
+    sentry_init, capture_envelopes, teardown_profiling
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        profile_lifecycle="trace",
+        profiler_mode="thread",
+        profile_session_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+            "continuous_profiling_auto_start": True,
+        },
+    )
+    capture_envelopes()
+
+    with sentry_sdk.traces.start_span(name="segment") as span:
+        time.sleep(0.1)
+        assert span._continuous_profile is not None
+        assert span._continuous_profile.active is True
+
+    assert span._continuous_profile.active is False
+
+    time.sleep(0.3)
+    assert get_profiler_id() is None, "profiler should have stopped"
+
+
 def test_transport_format(sentry_init, capture_envelopes):
     sentry_init(
         server_name="test-server",