Merge pull request #71 from modern-python/pyroscope

lesnik512 · web-flow · commit 83a843b7c9c3 · 2026-04-08T12:43:13.000+03:00
add pyroscope instrument
diff --git a/CLAUDE.md b/CLAUDE.md
@@ -61,6 +61,7 @@ Install via `pip install lite-bootstrap[<group>]` or `uv add lite-bootstrap[<gro
 | `litestar-all` | litestar + sentry + otl + logging |
 | `faststream-all` | faststream + sentry + otl + logging |
 | `free-all` | sentry + otl + logging |
+| `pyroscope` | pyroscope-io (add to any group) |
 
 ## Code style
 
diff --git a/README.md b/README.md
@@ -11,6 +11,7 @@ With `lite-bootstrap`, you receive an application with lightweight built-in supp
 - `sentry`
 - `prometheus`
 - `opentelemetry`
+- `pyroscope` - with OpenTelemetry trace-profile linking
 - `structlog`
 - `cors`
 - `swagger` - with additional offline version support
diff --git a/docs/index.md b/docs/index.md
@@ -8,6 +8,7 @@ With `lite-bootstrap`, you receive an application with lightweight built-in supp
 - `sentry`
 - `prometheus`
 - `opentelemetry`
+- `pyroscope` - with OpenTelemetry trace-profile linking
 - `structlog`
 - `cors`
 - `swagger` - with additional offline version support
diff --git a/docs/introduction/configuration.md b/docs/introduction/configuration.md
@@ -87,6 +87,22 @@ For FastStream you must provide additionally:
 - `opentelemetry_middleware_cls`
 
 
+## Pyroscope
+
+Pyroscope integration uses `pyroscope-io` package under the hood. Install it with the `pyroscope` extra, e.g. `lite-bootstrap[fastapi-all,pyroscope]`.
+
+To bootstrap Pyroscope, you must provide at least:
+
+- `pyroscope_endpoint` - the Pyroscope server address (e.g. `http://pyroscope:4040`).
+
+Additional parameters:
+
+- `pyroscope_sample_rate` - CPU profiling sample rate in Hz (default: `100`).
+- `pyroscope_tags` - key/value string pairs attached to all profiles.
+- `pyroscope_additional_params` - additional params passed directly to `pyroscope.configure`.
+
+When OpenTelemetry is also enabled, a `PyroscopeSpanProcessor` is automatically added to the tracer provider. It tags root spans with a `pyroscope.profile.id` attribute and sets Pyroscope thread tags so that traces and profiles can be linked in the Grafana UI.
+
 ## Structlog
 
 To bootstrap Structlog, you must set `service_debug` to False
diff --git a/docs/introduction/installation.md b/docs/introduction/installation.md
@@ -9,6 +9,7 @@ You can choose required framework and instruments using this table:
 | sentry        | `litestar-sentry`  | `faststream-sentry`  | `fastapi-sentry`  | `sentry`                             |
 | prometheus    | `litestar-metrics` | `faststream-metrics` | `fastapi-metrics` | not used                             |
 | opentelemetry | `litestar-otl`     | `faststream-otl`     | `fastapi-otl`     | `otl`                                |
+| pyroscope     | `pyroscope`        | `pyroscope`          | `pyroscope`       | `pyroscope`                          |
 | structlog     | `litestar-logging` | `faststream-logging` | `fastapi-logging` | `logging`                            |
 | cors          | no extra           | not used             | no extra          | not used                             |
 | swagger       | no extra           | not used             | no extra          | not used                             |
diff --git a/lite_bootstrap/__init__.py b/lite_bootstrap/__init__.py
@@ -2,6 +2,7 @@
 from lite_bootstrap.bootstrappers.faststream_bootstrapper import FastStreamBootstrapper, FastStreamConfig
 from lite_bootstrap.bootstrappers.free_bootstrapper import FreeBootstrapper, FreeBootstrapperConfig
 from lite_bootstrap.bootstrappers.litestar_bootstrapper import LitestarBootstrapper, LitestarConfig
+from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 
 
 __all__ = [
@@ -13,4 +14,6 @@
     "FreeBootstrapperConfig",
     "LitestarBootstrapper",
     "LitestarConfig",
+    "PyroscopeConfig",
+    "PyroscopeInstrument",
 ]
diff --git a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py
@@ -15,6 +15,7 @@
 from lite_bootstrap.instruments.logging_instrument import LoggingConfig, LoggingInstrument
 from lite_bootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpenTelemetryInstrument
 from lite_bootstrap.instruments.prometheus_instrument import PrometheusConfig, PrometheusInstrument
+from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 from lite_bootstrap.instruments.sentry_instrument import SentryConfig, SentryInstrument
 from lite_bootstrap.instruments.swagger_instrument import SwaggerConfig, SwaggerInstrument
 
@@ -36,7 +37,14 @@
 
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class FastAPIConfig(
-    CorsConfig, HealthChecksConfig, LoggingConfig, OpentelemetryConfig, PrometheusConfig, SentryConfig, SwaggerConfig
+    CorsConfig,
+    HealthChecksConfig,
+    LoggingConfig,
+    OpentelemetryConfig,
+    PrometheusConfig,
+    PyroscopeConfig,
+    SentryConfig,
+    SwaggerConfig,
 ):
     application: "fastapi.FastAPI" = dataclasses.field(default=None)  # ty: ignore[invalid-assignment]
     application_kwargs: dict[str, typing.Any] = dataclasses.field(default_factory=dict)
@@ -174,6 +182,7 @@ class FastAPIBootstrapper(BaseBootstrapper["fastapi.FastAPI"]):
     instruments_types: typing.ClassVar = [
         FastAPICorsInstrument,
         FastAPIOpenTelemetryInstrument,
+        PyroscopeInstrument,
         FastAPISentryInstrument,
         FastAPIHealthChecksInstrument,
         FastAPILoggingInstrument,
diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py
@@ -8,6 +8,7 @@
 from lite_bootstrap.instruments.logging_instrument import LoggingConfig, LoggingInstrument
 from lite_bootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpenTelemetryInstrument
 from lite_bootstrap.instruments.prometheus_instrument import PrometheusConfig, PrometheusInstrument
+from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 from lite_bootstrap.instruments.sentry_instrument import SentryConfig, SentryInstrument
 
 
@@ -51,7 +52,9 @@ def __init__(
 
 
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
-class FastStreamConfig(HealthChecksConfig, LoggingConfig, OpentelemetryConfig, PrometheusConfig, SentryConfig):
+class FastStreamConfig(
+    HealthChecksConfig, LoggingConfig, OpentelemetryConfig, PrometheusConfig, PyroscopeConfig, SentryConfig
+):
     application: "AsgiFastStream" = dataclasses.field(default_factory=AsgiFastStream)
     opentelemetry_middleware_cls: type[FastStreamTelemetryMiddlewareProtocol] | None = None
     prometheus_middleware_cls: type[FastStreamPrometheusMiddlewareProtocol] | None = None
@@ -152,6 +155,7 @@ class FastStreamBootstrapper(BaseBootstrapper["AsgiFastStream"]):
 
     instruments_types: typing.ClassVar = [
         FastStreamOpenTelemetryInstrument,
+        PyroscopeInstrument,
         FastStreamSentryInstrument,
         FastStreamHealthChecksInstrument,
         FastStreamLoggingInstrument,
diff --git a/lite_bootstrap/bootstrappers/free_bootstrapper.py b/lite_bootstrap/bootstrappers/free_bootstrapper.py
@@ -4,11 +4,12 @@
 from lite_bootstrap.bootstrappers.base import BaseBootstrapper
 from lite_bootstrap.instruments.logging_instrument import LoggingConfig, LoggingInstrument
 from lite_bootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpenTelemetryInstrument
+from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 from lite_bootstrap.instruments.sentry_instrument import SentryConfig, SentryInstrument
 
 
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
-class FreeBootstrapperConfig(LoggingConfig, OpentelemetryConfig, SentryConfig): ...
+class FreeBootstrapperConfig(LoggingConfig, OpentelemetryConfig, PyroscopeConfig, SentryConfig): ...
 
 
 class FreeBootstrapper(BaseBootstrapper[None]):
@@ -18,6 +19,7 @@ class FreeBootstrapper(BaseBootstrapper[None]):
         LoggingInstrument,
         SentryInstrument,
         OpenTelemetryInstrument,
+        PyroscopeInstrument,
     ]
     bootstrap_config: FreeBootstrapperConfig
     not_ready_message = ""
diff --git a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py
@@ -19,6 +19,7 @@
 from lite_bootstrap.instruments.prometheus_instrument import (
     PrometheusInstrument,
 )
+from lite_bootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 from lite_bootstrap.instruments.sentry_instrument import SentryConfig, SentryInstrument
 from lite_bootstrap.instruments.swagger_instrument import SwaggerConfig, SwaggerInstrument
 
@@ -83,7 +84,7 @@ async def handle(
                 default_span_details=build_litestar_route_details_from_scope,
                 excluded_urls=self._excluded_urls,
                 tracer_provider=self._tracer_provider,
-            )(scope, receive, send)  # ty: ignore
+            )(scope, receive, send)  # ty: ignore[invalid-argument-type]
 
 
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
@@ -93,6 +94,7 @@ class LitestarConfig(
     LoggingConfig,
     OpentelemetryConfig,
     PrometheusBootstrapperConfig,
+    PyroscopeConfig,
     SentryConfig,
     SwaggerConfig,
 ):
@@ -239,6 +241,7 @@ class LitestarBootstrapper(BaseBootstrapper["litestar.Litestar"]):
     instruments_types: typing.ClassVar = [
         LitestarCorsInstrument,
         LitestarOpenTelemetryInstrument,
+        PyroscopeInstrument,
         LitestarSentryInstrument,
         LitestarHealthChecksInstrument,
         LitestarLoggingInstrument,
diff --git a/lite_bootstrap/import_checker.py b/lite_bootstrap/import_checker.py
@@ -15,3 +15,4 @@
 is_litestar_opentelemetry_installed = (
     is_opentelemetry_installed and is_litestar_installed and find_spec("opentelemetry.instrumentation.asgi") is not None
 )
+is_pyroscope_installed = find_spec("pyroscope") is not None
diff --git a/lite_bootstrap/instruments/opentelemetry_instrument.py b/lite_bootstrap/instruments/opentelemetry_instrument.py
@@ -1,4 +1,5 @@
 import dataclasses
+import logging
 import os
 import typing
 
@@ -10,11 +11,19 @@
     from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 
 if import_checker.is_opentelemetry_installed:
+    from opentelemetry.context import Context
     from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
     from opentelemetry.sdk import resources
-    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
     from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor
-    from opentelemetry.trace import set_tracer_provider
+    from opentelemetry.trace import Span, format_span_id, set_tracer_provider
+
+if import_checker.is_pyroscope_installed:
+    import pyroscope
+
+
+def _format_span(readable_span: "ReadableSpan") -> 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
@@ -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
@@ -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
diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py
diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py
diff --git a/tests/test_free_bootstrap.py b/tests/test_free_bootstrap.py
