diff --git a/sentry_sdk/traces.py b/sentry_sdk/traces.py index e10f2ca2a7..944e17e5d7 100644 --- a/sentry_sdk/traces.py +++ b/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__ = ( diff --git a/tests/tracing/test_span_streaming.py b/tests/tracing/test_span_streaming.py index cea83be7ec..8b9b658c77 100644 --- a/tests/tracing/test_span_streaming.py +++ b/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",