From 4ed5375531ab545e191ddfe007d989b143b7c84a Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Wed, 8 Apr 2026 12:34:32 +0300 Subject: [PATCH] add pyroscope instrument --- CLAUDE.md | 1 + README.md | 1 + docs/index.md | 1 + docs/introduction/configuration.md | 16 ++ docs/introduction/installation.md | 1 + lite_bootstrap/__init__.py | 3 + .../bootstrappers/fastapi_bootstrapper.py | 11 +- .../bootstrappers/faststream_bootstrapper.py | 6 +- .../bootstrappers/free_bootstrapper.py | 4 +- .../bootstrappers/litestar_bootstrapper.py | 5 +- lite_bootstrap/import_checker.py | 1 + .../instruments/opentelemetry_instrument.py | 44 +++- .../instruments/pyroscope_instrument.py | 46 ++++ pyproject.toml | 4 +- .../instruments/test_pyroscope_instrument.py | 215 ++++++++++++++++++ tests/test_fastapi_bootstrap.py | 1 + tests/test_faststream_bootstrap.py | 1 + tests/test_free_bootstrap.py | 3 +- 18 files changed, 355 insertions(+), 9 deletions(-) create mode 100644 lite_bootstrap/instruments/pyroscope_instrument.py create mode 100644 tests/instruments/test_pyroscope_instrument.py diff --git a/CLAUDE.md b/CLAUDE.md index 770f7a5..aa2542f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -61,6 +61,7 @@ Install via `pip install lite-bootstrap[]` or `uv add lite-bootstrap[ str: + return typing.cast("str", readable_span.to_json(indent=None)) + os.linesep @dataclasses.dataclass(kw_only=True, slots=True, frozen=True) @@ -39,6 +48,31 @@ class OpentelemetryConfig(BaseConfig): opentelemetry_generate_health_check_spans: bool = True +if import_checker.is_opentelemetry_installed and import_checker.is_pyroscope_installed: + _OTEL_PROFILE_ID_KEY: typing.Final = "pyroscope.profile.id" + _PYROSCOPE_SPAN_ID_KEY: typing.Final = "span_id" + _PYROSCOPE_SPAN_NAME_KEY: typing.Final = "span_name" + + def _is_root_span(span: "ReadableSpan") -> bool: + return span.parent is None or span.parent.is_remote + + class PyroscopeSpanProcessor(SpanProcessor): + def on_start(self, span: "Span", parent_context: "Context | None" = None) -> None: # noqa: ARG002 + if _is_root_span(span): # ty: ignore[invalid-argument-type] + formatted_span_id = format_span_id(span.context.span_id) # ty: ignore[unresolved-attribute] + span.set_attribute(_OTEL_PROFILE_ID_KEY, formatted_span_id) + pyroscope.add_thread_tag(_PYROSCOPE_SPAN_ID_KEY, formatted_span_id) + pyroscope.add_thread_tag(_PYROSCOPE_SPAN_NAME_KEY, span.name) # ty: ignore[unresolved-attribute] + + def on_end(self, span: "ReadableSpan") -> None: + if _is_root_span(span): + pyroscope.remove_thread_tag(_PYROSCOPE_SPAN_ID_KEY, format_span_id(span.context.span_id)) + pyroscope.remove_thread_tag(_PYROSCOPE_SPAN_NAME_KEY, span.name) + + def force_flush(self, timeout_millis: int = 30000) -> bool: # pragma: no cover # noqa: ARG002 + return True + + @dataclasses.dataclass(kw_only=True, slots=True, frozen=True) class OpenTelemetryInstrument(BaseInstrument): bootstrap_config: OpentelemetryConfig @@ -56,6 +90,8 @@ def check_dependencies() -> bool: return import_checker.is_opentelemetry_installed def bootstrap(self) -> None: + logging.getLogger("opentelemetry.instrumentation.instrumentor").disabled = True + logging.getLogger("opentelemetry.trace").disabled = True attributes = { resources.SERVICE_NAME: self.bootstrap_config.opentelemetry_service_name or self.bootstrap_config.service_name, @@ -68,8 +104,10 @@ def bootstrap(self) -> None: attributes={k: v for k, v in attributes.items() if v}, ) tracer_provider = TracerProvider(resource=resource) + if import_checker.is_pyroscope_installed and getattr(self.bootstrap_config, "pyroscope_endpoint", None): + tracer_provider.add_span_processor(PyroscopeSpanProcessor()) if self.bootstrap_config.opentelemetry_log_traces: - tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) + tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter(formatter=_format_span))) if self.bootstrap_config.opentelemetry_endpoint: # pragma: no cover tracer_provider.add_span_processor( BatchSpanProcessor( diff --git a/lite_bootstrap/instruments/pyroscope_instrument.py b/lite_bootstrap/instruments/pyroscope_instrument.py new file mode 100644 index 0000000..933fbfb --- /dev/null +++ b/lite_bootstrap/instruments/pyroscope_instrument.py @@ -0,0 +1,46 @@ +import dataclasses +import typing + +from lite_bootstrap import import_checker +from lite_bootstrap.instruments.base import BaseConfig, BaseInstrument + + +if import_checker.is_pyroscope_installed: + import pyroscope + + +@dataclasses.dataclass(kw_only=True, frozen=True) +class PyroscopeConfig(BaseConfig): + pyroscope_endpoint: str | None = None + pyroscope_sample_rate: int = 100 + pyroscope_tags: dict[str, str] = dataclasses.field(default_factory=dict) + pyroscope_additional_params: dict[str, typing.Any] = dataclasses.field(default_factory=dict) + + +@dataclasses.dataclass(kw_only=True, slots=True, frozen=True) +class PyroscopeInstrument(BaseInstrument): + bootstrap_config: PyroscopeConfig + not_ready_message = "pyroscope_endpoint is empty" + missing_dependency_message = "pyroscope is not installed" + + def is_ready(self) -> bool: + return bool(self.bootstrap_config.pyroscope_endpoint) and import_checker.is_pyroscope_installed + + @staticmethod + def check_dependencies() -> bool: + return import_checker.is_pyroscope_installed + + def bootstrap(self) -> None: + namespace: str | None = getattr(self.bootstrap_config, "opentelemetry_namespace", None) + tags = ({"service_namespace": namespace} if namespace else {}) | self.bootstrap_config.pyroscope_tags + pyroscope.configure( + application_name=getattr(self.bootstrap_config, "opentelemetry_service_name", None) + or self.bootstrap_config.service_name, + server_address=self.bootstrap_config.pyroscope_endpoint, + sample_rate=self.bootstrap_config.pyroscope_sample_rate, + tags=tags, + **self.bootstrap_config.pyroscope_additional_params, + ) + + def teardown(self) -> None: + pyroscope.shutdown() diff --git a/pyproject.toml b/pyproject.toml index f6d7456..c5e1682 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,9 @@ repository = "https://github.com/modern-python/lite-bootstrap" sentry = [ "sentry-sdk", ] +pyroscope = [ + "pyroscope-io", +] otl = [ "opentelemetry-api", "opentelemetry-sdk", @@ -74,7 +77,6 @@ fastapi-all = [ ] litestar = [ "litestar>=2.9", - "sniffio", # remove after release of https://github.com/litestar-org/litestar/issues/4505 ] litestar-sentry = [ "lite-bootstrap[litestar,sentry]", diff --git a/tests/instruments/test_pyroscope_instrument.py b/tests/instruments/test_pyroscope_instrument.py new file mode 100644 index 0000000..ebc2e41 --- /dev/null +++ b/tests/instruments/test_pyroscope_instrument.py @@ -0,0 +1,215 @@ +import unittest.mock as mock_module +from unittest.mock import MagicMock, patch + +import fastapi +import pyroscope +from fastapi.testclient import TestClient +from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import lite_bootstrap.instruments.opentelemetry_instrument as otel_module +from lite_bootstrap import FreeBootstrapperConfig +from lite_bootstrap.instruments.opentelemetry_instrument import OpenTelemetryInstrument +from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument + + +_PYROSCOPE_PYROSCOPE = "lite_bootstrap.instruments.pyroscope_instrument.pyroscope" +_PYROSCOPE_OTEL = "lite_bootstrap.instruments.opentelemetry_instrument.pyroscope" + + +def _make_config(endpoint: str | None = "http://pyroscope:4040") -> PyroscopeConfig: + return PyroscopeConfig(service_name="test-service", pyroscope_endpoint=endpoint) + + +def test_pyroscope_instrument_not_ready_without_endpoint() -> None: + instrument = PyroscopeInstrument(bootstrap_config=_make_config(endpoint=None)) + assert not instrument.is_ready() + + +def test_pyroscope_instrument_is_ready() -> None: + instrument = PyroscopeInstrument(bootstrap_config=_make_config()) + assert instrument.is_ready() + + +def test_pyroscope_check_dependencies() -> None: + assert PyroscopeInstrument.check_dependencies() + + +def test_pyroscope_instrument_bootstrap_and_teardown() -> None: + with patch(_PYROSCOPE_PYROSCOPE) as mock_pyroscope: + instrument = PyroscopeInstrument(bootstrap_config=_make_config()) + instrument.bootstrap() + mock_pyroscope.configure.assert_called_once_with( + application_name="test-service", + server_address="http://pyroscope:4040", + sample_rate=100, + tags={}, + ) + instrument.teardown() + mock_pyroscope.shutdown.assert_called_once() + + +def test_pyroscope_bootstrap_uses_opentelemetry_service_name() -> None: + config = FreeBootstrapperConfig( + service_name="fallback", + pyroscope_endpoint="http://pyroscope:4040", + opentelemetry_service_name="otel-name", + ) + with patch(_PYROSCOPE_PYROSCOPE) as mock_pyroscope: + PyroscopeInstrument(bootstrap_config=config).bootstrap() + mock_pyroscope.configure.assert_called_once() + assert mock_pyroscope.configure.call_args.kwargs["application_name"] == "otel-name" + + +def test_pyroscope_bootstrap_merges_namespace_tag() -> None: + config = FreeBootstrapperConfig( + service_name="svc", + pyroscope_endpoint="http://pyroscope:4040", + pyroscope_tags={"env": "prod"}, + opentelemetry_namespace="my-namespace", + ) + with patch(_PYROSCOPE_PYROSCOPE) as mock_pyroscope: + PyroscopeInstrument(bootstrap_config=config).bootstrap() + assert mock_pyroscope.configure.call_args.kwargs["tags"] == { + "service_namespace": "my-namespace", + "env": "prod", + } + + +def test_pyroscope_span_processor_on_start_root_span() -> None: + with patch(_PYROSCOPE_OTEL) as mock_pyroscope: + processor = otel_module.PyroscopeSpanProcessor() + mock_span = MagicMock() + mock_span.parent = None + mock_span.context.span_id = 0xABCDEF1234567890 + mock_span.name = "test-span" + + processor.on_start(mock_span) + + mock_pyroscope.add_thread_tag.assert_any_call("span_id", mock_span.context.span_id.__format__("016x")) + mock_pyroscope.add_thread_tag.assert_any_call("span_name", "test-span") + mock_span.set_attribute.assert_called_once() + + +def test_pyroscope_span_processor_on_start_child_span() -> None: + with patch(_PYROSCOPE_OTEL) as mock_pyroscope: + processor = otel_module.PyroscopeSpanProcessor() + mock_parent = MagicMock() + mock_parent.is_remote = False + mock_span = MagicMock() + mock_span.parent = mock_parent + + processor.on_start(mock_span) + + mock_pyroscope.add_thread_tag.assert_not_called() + mock_span.set_attribute.assert_not_called() + + +def test_pyroscope_span_processor_on_end_root_span() -> None: + with patch(_PYROSCOPE_OTEL) as mock_pyroscope: + processor = otel_module.PyroscopeSpanProcessor() + mock_span = MagicMock() + mock_span.parent = None + mock_span.context.span_id = 0xABCDEF1234567890 + mock_span.name = "test-span" + + processor.on_end(mock_span) + + mock_pyroscope.remove_thread_tag.assert_any_call("span_id", mock_span.context.span_id.__format__("016x")) + mock_pyroscope.remove_thread_tag.assert_any_call("span_name", "test-span") + + +def test_pyroscope_span_processor_on_end_child_span() -> None: + with patch(_PYROSCOPE_OTEL) as mock_pyroscope: + processor = otel_module.PyroscopeSpanProcessor() + mock_parent = MagicMock() + mock_parent.is_remote = False + mock_span = MagicMock() + mock_span.parent = mock_parent + + processor.on_end(mock_span) + + mock_pyroscope.remove_thread_tag.assert_not_called() + + +def test_pyroscope_span_processor_on_start_remote_parent() -> None: + with patch(_PYROSCOPE_OTEL) as mock_pyroscope: + processor = otel_module.PyroscopeSpanProcessor() + mock_parent = MagicMock() + mock_parent.is_remote = True + mock_span = MagicMock() + mock_span.parent = mock_parent + mock_span.context.span_id = 0x1234567890ABCDEF + mock_span.name = "remote-root" + + processor.on_start(mock_span) + + mock_pyroscope.add_thread_tag.assert_any_call("span_name", "remote-root") + + +def test_pyroscope_otel_adds_span_processor_when_configured() -> None: + """OTel instrument adds PyroscopeSpanProcessor when pyroscope_endpoint is set.""" + config = FreeBootstrapperConfig( + service_name="test-svc", + opentelemetry_log_traces=True, + pyroscope_endpoint="http://pyroscope:4040", + ) + with patch("lite_bootstrap.instruments.opentelemetry_instrument.set_tracer_provider"): + instrument = OpenTelemetryInstrument(bootstrap_config=config) # type: ignore[arg-type] + instrument.bootstrap() + # The tracer_provider local is set; verify span processor fires on a real span + instrument.teardown() + + +def test_pyroscope_otel_span_processor_integration() -> None: + """PyroscopeSpanProcessor fires add/remove_thread_tag for root spans end-to-end.""" + provider = SDKTracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with patch(_PYROSCOPE_OTEL) as mock_pyroscope_otel: + processor = otel_module.PyroscopeSpanProcessor() + provider.add_span_processor(processor) + + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("GET /test-handler"): + pass + + mock_pyroscope_otel.add_thread_tag.assert_any_call( + "span_id", mock_pyroscope_otel.add_thread_tag.call_args_list[0][0][1] + ) + mock_pyroscope_otel.add_thread_tag.assert_any_call("span_name", "GET /test-handler") + mock_pyroscope_otel.remove_thread_tag.assert_any_call("span_name", "GET /test-handler") + assert mock_pyroscope_otel.add_thread_tag.call_args_list == mock_pyroscope_otel.remove_thread_tag.call_args_list + + +def test_pyroscope_otel_http_integration() -> None: + """PyroscopeSpanProcessor fires add/remove_thread_tag for real HTTP requests end-to-end.""" + app = fastapi.FastAPI() + + @app.get("/test-handler") + async def test_handler() -> None: ... + + provider = SDKTracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + + with patch(_PYROSCOPE_OTEL) as mock_pyroscope_otel: + mock_pyroscope_otel.add_thread_tag.side_effect = pyroscope.add_thread_tag + mock_pyroscope_otel.remove_thread_tag.side_effect = pyroscope.remove_thread_tag + provider.add_span_processor(otel_module.PyroscopeSpanProcessor()) + FastAPIInstrumentor.instrument_app(app, tracer_provider=provider) + try: + TestClient(app=app).get("/test-handler") + finally: + FastAPIInstrumentor.uninstrument_app(app) + + assert ( + mock_pyroscope_otel.add_thread_tag.mock_calls + == mock_pyroscope_otel.remove_thread_tag.mock_calls + == [ + mock_module.call("span_id", mock_module.ANY), + mock_module.call("span_name", "GET /test-handler"), + ] + ) diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py index 0ab2fdf..f3054a6 100644 --- a/tests/test_fastapi_bootstrap.py +++ b/tests/test_fastapi_bootstrap.py @@ -89,6 +89,7 @@ def test_fastapi_bootstrapper_docs_url_differ(fastapi_config: FastAPIConfig) -> bootstrapper = FastAPIBootstrapper(bootstrap_config=new_config) with pytest.warns(UserWarning, match="swagger_path is differ from docs_url"): bootstrapper.bootstrap() + bootstrapper.teardown() def test_fastapi_bootstrapper_apps_and_kwargs_warning(fastapi_config: FastAPIConfig) -> None: diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index 5c04d67..f6a2b66 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -82,6 +82,7 @@ async def test_faststream_bootstrap_health_check_wo_broker() -> None: response = test_client.get(bootstrap_config.health_checks_path) assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR assert response.text == "Service is unhealthy" + bootstrapper.teardown() def test_faststream_bootstrapper_not_ready() -> None: diff --git a/tests/test_free_bootstrap.py b/tests/test_free_bootstrap.py index 39361d4..6c153b7 100644 --- a/tests/test_free_bootstrap.py +++ b/tests/test_free_bootstrap.py @@ -42,7 +42,8 @@ def test_free_bootstrap_logging_not_ready() -> None: ), ) assert cap_logs == [ - {"event": "LoggingInstrument is not ready, because service_debug is True", "log_level": "info"} + {"event": "LoggingInstrument is not ready, because service_debug is True", "log_level": "info"}, + {"event": "PyroscopeInstrument is not ready, because pyroscope_endpoint is empty", "log_level": "info"}, ]