test_continuous_profiler.py

import threading
import time
from collections import defaultdict
from unittest import mock

import pytest

import sentry_sdk
from sentry_sdk.consts import VERSION
from sentry_sdk.profiler.continuous_profiler import (
    is_profile_session_sampled,
    get_profiler_id,
    setup_continuous_profiler,
    start_profiler,
    start_profile_session,
    stop_profiler,
    stop_profile_session,
)
from tests.conftest import ApproxDict

try:
    import gevent
except ImportError:
    gevent = None


requires_gevent = pytest.mark.skipif(gevent is None, reason="gevent not enabled")


def get_client_options(use_top_level_profiler_mode):
    def client_options(
        mode=None, auto_start=None, profile_session_sample_rate=1.0, lifecycle="manual"
    ):
        if use_top_level_profiler_mode:
            return {
                "profile_lifecycle": lifecycle,
                "profiler_mode": mode,
                "profile_session_sample_rate": profile_session_sample_rate,
                "_experiments": {
                    "continuous_profiling_auto_start": auto_start,
                },
            }
        return {
            "profile_lifecycle": lifecycle,
            "profile_session_sample_rate": profile_session_sample_rate,
            "_experiments": {
                "continuous_profiling_auto_start": auto_start,
                "continuous_profiling_mode": mode,
            },
        }

    return client_options


mock_sdk_info = {
    "name": "sentry.python",
    "version": VERSION,
    "packages": [{"name": "pypi:sentry-sdk", "version": VERSION}],
}


@pytest.mark.parametrize("mode", [pytest.param("foo")])
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
def test_continuous_profiler_invalid_mode(mode, make_options, teardown_profiling):
    with pytest.raises(ValueError):
        setup_continuous_profiler(
            make_options(mode=mode),
            mock_sdk_info,
            lambda envelope: None,
        )


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
def test_continuous_profiler_valid_mode(mode, make_options, teardown_profiling):
    options = make_options(mode=mode)
    setup_continuous_profiler(
        options,
        mock_sdk_info,
        lambda envelope: None,
    )


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
def test_continuous_profiler_setup_twice(mode, make_options, teardown_profiling):
    assert not is_profile_session_sampled()

    # setting up the first time should return True to indicate success
    options = make_options(mode=mode, profile_session_sample_rate=1.0)
    assert setup_continuous_profiler(
        options,
        mock_sdk_info,
        lambda envelope: None,
    )
    assert is_profile_session_sampled()

    # setting up the second time should return True to indicate re-init
    options = make_options(mode=mode, profile_session_sample_rate=0.0)
    assert setup_continuous_profiler(
        options,
        mock_sdk_info,
        lambda envelope: None,
    )
    assert not is_profile_session_sampled()


def assert_single_transaction_with_profile_chunks(
    envelopes, thread, max_chunks=None, transactions=1
):
    items = defaultdict(list)
    for envelope in envelopes:
        for item in envelope.items:
            items[item.type].append(item)

    assert len(items["transaction"]) == transactions
    assert len(items["profile_chunk"]) > 0
    if max_chunks is not None:
        assert len(items["profile_chunk"]) <= max_chunks

    for chunk_item in items["profile_chunk"]:
        chunk = chunk_item.payload.json
        headers = chunk_item.headers
        assert chunk["platform"] == headers["platform"]

    transaction = items["transaction"][0].payload.json

    trace_context = transaction["contexts"]["trace"]

    assert trace_context == ApproxDict(
        {
            "data": ApproxDict(
                {
                    "thread.id": str(thread.ident),
                    "thread.name": thread.name,
                }
            ),
        }
    )

    profile_context = transaction["contexts"]["profile"]
    profiler_id = profile_context["profiler_id"]

    assert profile_context == ApproxDict({"profiler_id": profiler_id})

    spans = transaction["spans"]
    assert len(spans) > 0
    for span in spans:
        assert span["data"] == ApproxDict(
            {
                "profiler_id": profiler_id,
                "thread.id": str(thread.ident),
                "thread.name": thread.name,
            }
        )

    for profile_chunk_item in items["profile_chunk"]:
        profile_chunk = profile_chunk_item.payload.json
        del profile_chunk["profile"]  # make the diff easier to read
        assert profile_chunk == ApproxDict(
            {
                "client_sdk": {
                    "name": mock.ANY,
                    "version": VERSION,
                },
                "platform": "python",
                "profiler_id": profiler_id,
                "version": "2",
            }
        )


def assert_single_transaction_without_profile_chunks(envelopes):
    items = defaultdict(list)
    for envelope in envelopes:
        for item in envelope.items:
            items[item.type].append(item)

    assert len(items["transaction"]) == 1
    assert len(items["profile_chunk"]) == 0

    transaction = items["transaction"][0].payload.json
    assert "profile" not in transaction["contexts"]


@pytest.mark.forked
@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    ["start_profiler_func", "stop_profiler_func"],
    [
        pytest.param(
            start_profile_session,
            stop_profile_session,
            id="start_profile_session/stop_profile_session (deprecated)",
        ),
        pytest.param(
            start_profiler,
            stop_profiler,
            id="start_profiler/stop_profiler",
        ),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
@mock.patch("sentry_sdk.profiler.continuous_profiler.PROFILE_BUFFER_SECONDS", 0.01)
def test_continuous_profiler_auto_start_and_manual_stop(
    sentry_init,
    capture_envelopes,
    mode,
    start_profiler_func,
    stop_profiler_func,
    make_options,
    teardown_profiling,
):
    options = make_options(mode=mode, auto_start=True)
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    envelopes = capture_envelopes()

    thread = threading.current_thread()

    with sentry_sdk.start_transaction(name="profiling"):
        with sentry_sdk.start_span(op="op"):
            time.sleep(0.05)

    assert_single_transaction_with_profile_chunks(envelopes, thread)

    for _ in range(3):
        stop_profiler_func()

        envelopes.clear()

        with sentry_sdk.start_transaction(name="profiling"):
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.05)

        assert_single_transaction_without_profile_chunks(envelopes)

        start_profiler_func()

        envelopes.clear()

        with sentry_sdk.start_transaction(name="profiling"):
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.05)

        assert_single_transaction_with_profile_chunks(envelopes, thread)


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    ["start_profiler_func", "stop_profiler_func"],
    [
        pytest.param(
            start_profile_session,
            stop_profile_session,
            id="start_profile_session/stop_profile_session  (deprecated)",
        ),
        pytest.param(
            start_profiler,
            stop_profiler,
            id="start_profiler/stop_profiler",
        ),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
@mock.patch("sentry_sdk.profiler.continuous_profiler.PROFILE_BUFFER_SECONDS", 0.01)
def test_continuous_profiler_manual_start_and_stop_sampled(
    sentry_init,
    capture_envelopes,
    mode,
    start_profiler_func,
    stop_profiler_func,
    make_options,
    teardown_profiling,
):
    options = make_options(
        mode=mode, profile_session_sample_rate=1.0, lifecycle="manual"
    )
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    envelopes = capture_envelopes()

    thread = threading.current_thread()

    for _ in range(3):
        start_profiler_func()

        envelopes.clear()

        with sentry_sdk.start_transaction(name="profiling"):
            assert get_profiler_id() is not None, "profiler should be running"
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.1)
            assert get_profiler_id() is not None, "profiler should be running"

        assert_single_transaction_with_profile_chunks(envelopes, thread)

        assert get_profiler_id() is not None, "profiler should be running"

        stop_profiler_func()

        # the profiler stops immediately in manual mode
        assert get_profiler_id() is None, "profiler should not be running"

        envelopes.clear()

        with sentry_sdk.start_transaction(name="profiling"):
            assert get_profiler_id() is None, "profiler should not be running"
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.1)
            assert get_profiler_id() is None, "profiler should not be running"

        assert_single_transaction_without_profile_chunks(envelopes)


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    ["start_profiler_func", "stop_profiler_func"],
    [
        pytest.param(
            start_profile_session,
            stop_profile_session,
            id="start_profile_session/stop_profile_session (deprecated)",
        ),
        pytest.param(
            start_profiler,
            stop_profiler,
            id="start_profiler/stop_profiler",
        ),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
def test_continuous_profiler_manual_start_and_stop_unsampled(
    sentry_init,
    capture_envelopes,
    mode,
    start_profiler_func,
    stop_profiler_func,
    make_options,
    teardown_profiling,
):
    options = make_options(
        mode=mode, profile_session_sample_rate=0.0, lifecycle="manual"
    )
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    envelopes = capture_envelopes()

    start_profiler_func()

    with sentry_sdk.start_transaction(name="profiling"):
        with sentry_sdk.start_span(op="op"):
            time.sleep(0.05)

    assert_single_transaction_without_profile_chunks(envelopes)

    stop_profiler_func()


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
@mock.patch("sentry_sdk.profiler.continuous_profiler.DEFAULT_SAMPLING_FREQUENCY", 21)
def test_continuous_profiler_auto_start_and_stop_sampled(
    sentry_init,
    capture_envelopes,
    mode,
    make_options,
    teardown_profiling,
):
    options = make_options(
        mode=mode, profile_session_sample_rate=1.0, lifecycle="trace"
    )
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    envelopes = capture_envelopes()

    thread = threading.current_thread()

    all_profiler_ids = set()

    for _ in range(3):
        envelopes.clear()

        profiler_ids = set()

        with sentry_sdk.start_transaction(name="profiling 1"):
            profiler_id = get_profiler_id()
            assert profiler_id is not None, "profiler should be running"
            profiler_ids.add(profiler_id)
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.1)
            profiler_id = get_profiler_id()
            assert profiler_id is not None, "profiler should be running"
            profiler_ids.add(profiler_id)

        time.sleep(0.03)

        # the profiler takes a while to stop in auto mode so if we start
        # a transaction immediately, it'll be part of the same chunk
        profiler_id = get_profiler_id()
        assert profiler_id is not None, "profiler should be running"
        profiler_ids.add(profiler_id)

        with sentry_sdk.start_transaction(name="profiling 2"):
            profiler_id = get_profiler_id()
            assert profiler_id is not None, "profiler should be running"
            profiler_ids.add(profiler_id)
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.1)
            profiler_id = get_profiler_id()
            assert profiler_id is not None, "profiler should be running"
            profiler_ids.add(profiler_id)

        # wait at least 1 cycle for the profiler to stop
        time.sleep(0.2)
        assert get_profiler_id() is None, "profiler should not be running"

        assert len(profiler_ids) == 1
        all_profiler_ids.add(profiler_ids.pop())

        assert_single_transaction_with_profile_chunks(
            envelopes, thread, max_chunks=1, transactions=2
        )

    assert len(all_profiler_ids) == 3


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param("thread"),
        pytest.param("gevent", marks=requires_gevent),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
@mock.patch("sentry_sdk.profiler.continuous_profiler.PROFILE_BUFFER_SECONDS", 0.01)
def test_continuous_profiler_auto_start_and_stop_unsampled(
    sentry_init,
    capture_envelopes,
    mode,
    make_options,
    teardown_profiling,
):
    options = make_options(
        mode=mode, profile_session_sample_rate=0.0, lifecycle="trace"
    )
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    envelopes = capture_envelopes()

    for _ in range(3):
        envelopes.clear()

        with sentry_sdk.start_transaction(name="profiling"):
            assert get_profiler_id() is None, "profiler should not be running"
            with sentry_sdk.start_span(op="op"):
                time.sleep(0.05)
            assert get_profiler_id() is None, "profiler should not be running"

        assert get_profiler_id() is None, "profiler should not be running"
        assert_single_transaction_without_profile_chunks(envelopes)


@pytest.mark.parametrize(
    ["mode", "class_name"],
    [
        pytest.param("thread", "ThreadContinuousScheduler"),
        pytest.param(
            "gevent",
            "GeventContinuousScheduler",
            marks=requires_gevent,
        ),
    ],
)
@pytest.mark.parametrize(
    ["start_profiler_func", "stop_profiler_func"],
    [
        pytest.param(
            start_profile_session,
            stop_profile_session,
            id="start_profile_session/stop_profile_session (deprecated)",
        ),
        pytest.param(
            start_profiler,
            stop_profiler,
            id="start_profiler/stop_profiler",
        ),
    ],
)
@pytest.mark.parametrize(
    "make_options",
    [
        pytest.param(get_client_options(True), id="non-experiment"),
        pytest.param(get_client_options(False), id="experiment"),
    ],
)
def test_continuous_profiler_manual_start_and_stop_noop_when_using_trace_lifecyle(
    sentry_init,
    mode,
    start_profiler_func,
    stop_profiler_func,
    class_name,
    make_options,
    teardown_profiling,
):
    options = make_options(
        mode=mode, profile_session_sample_rate=0.0, lifecycle="trace"
    )
    sentry_init(
        traces_sample_rate=1.0,
        **options,
    )

    with mock.patch(
        f"sentry_sdk.profiler.continuous_profiler.{class_name}.ensure_running"
    ) as mock_ensure_running:
        start_profiler_func()
        mock_ensure_running.assert_not_called()

    with mock.patch(
        f"sentry_sdk.profiler.continuous_profiler.{class_name}.teardown"
    ) as mock_teardown:
        stop_profiler_func()
        mock_teardown.assert_not_called()


@mock.patch("sentry_sdk.profiler.continuous_profiler.PROFILE_BUFFER_SECONDS", 0.01)
def test_continuous_profiler_run_does_not_null_buffer(
    sentry_init,
    capture_envelopes,
    teardown_profiling,
):
    """
    Verifies that ContinuousScheduler.run() does not set self.buffer = None
    after exiting its sampling loop.

    Previously, run() would execute `self.buffer = None` after the while
    loop exited. During rapid stop/start cycles, this could race with
    ensure_running() which creates a new buffer: the old thread's cleanup
    would destroy the newly-created buffer, causing the new profiler thread
    to silently drop all samples (self.buffer is None in the sampler).

    The fix uses a local buffer reference for flushing and never sets
    self.buffer = None from run().
    """
    from sentry_sdk.profiler import continuous_profiler as cp
    from sentry_sdk.profiler.continuous_profiler import ContinuousScheduler

    options = get_client_options(True)(
        mode="thread", profile_session_sample_rate=1.0, lifecycle="manual"
    )
    sentry_init(traces_sample_rate=1.0, **options)
    envelopes = capture_envelopes()
    thread = threading.current_thread()

    # Start and verify profiler works
    start_profiler()
    envelopes.clear()
    with sentry_sdk.start_transaction(name="profiling"):
        with sentry_sdk.start_span(op="op"):
            time.sleep(0.1)
    assert_single_transaction_with_profile_chunks(envelopes, thread)

    # Get the scheduler and create a sentinel buffer.
    # We'll call run() directly to verify it doesn't null out self.buffer.
    scheduler = cp._scheduler
    assert scheduler is not None

    # Stop the profiler so the thread exits cleanly
    stop_profiler()

    # Now set up a fresh buffer and mark the scheduler as not running
    # (simulating the state right after ensure_running() created a new buffer
    # but the old thread hasn't done cleanup yet).
    scheduler.reset_buffer()
    buffer_before = scheduler.buffer
    assert buffer_before is not None

    # Simulate what happens when run() exits its while loop:
    # self.running is already False, so the while loop exits immediately.
    scheduler.running = False
    scheduler.run()

    # After the fix, run() should NOT have set self.buffer = None.
    # It should only flush using a local reference.
    assert scheduler.buffer is not None, (
        "run() must not set self.buffer = None; "
        "this would destroy buffers created by concurrent ensure_running() calls"
    )