Set span_name tag on Pyroscope profiles (#96)

* Enhance Pyroscope integration by customizing span processor and updating dependencies

* Add type ignore comment for untyped pyroscope import

* Adjust span name key and update test for Pyroscope instrumentation

Author: vrslev

microbootstrap/instruments/opentelemetry_instrument.py

@@ -1,28 +1,30 @@
 from __future__ import annotations
 import dataclasses
 import os
+import threading
 import typing
 
 import pydantic
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore[attr-defined] # noqa: TC002
 from opentelemetry.sdk import resources
-from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor
-from opentelemetry.trace import set_tracer_provider
+from opentelemetry.trace import format_span_id, set_tracer_provider
 
 from microbootstrap.instruments.base import BaseInstrumentConfig, Instrument
 
 
 try:
-    from pyroscope.otel import PyroscopeSpanProcessor  # type: ignore[import-untyped]
+    import pyroscope  # type: ignore[import-untyped]
 except ImportError:  # pragma: no cover
-    PyroscopeSpanProcessor = None
+    pyroscope = None
 
 
 if typing.TYPE_CHECKING:
     import faststream
+    from opentelemetry.context import Context
     from opentelemetry.metrics import Meter, MeterProvider
     from opentelemetry.trace import TracerProvider
 
@@ -97,7 +99,7 @@ def bootstrap(self) -> None:
         resource: typing.Final = resources.Resource.create(attributes=attributes)
 
         self.tracer_provider = SdkTracerProvider(resource=resource)
-        if self.instrument_config.pyroscope_endpoint and PyroscopeSpanProcessor:
+        if self.instrument_config.pyroscope_endpoint and pyroscope:
             self.tracer_provider.add_span_processor(PyroscopeSpanProcessor())
 
         if self.instrument_config.service_debug:
@@ -129,3 +131,34 @@ def define_exclude_urls(self) -> list[str]:
     @classmethod
     def get_config_type(cls) -> type[OpentelemetryConfig]:
         return OpentelemetryConfig
+
+
+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
+
+
+# Extended `pyroscope-otel` span processor: https://github.com/grafana/otel-profiling-python/blob/990662d416943e992ab70036b35b27488c98336a/src/pyroscope/otel/__init__.py
+# Includes `span_name` to identify if it makes sense to go to profiles from traces.
+class PyroscopeSpanProcessor(SpanProcessor):
+    def on_start(self, span: Span, parent_context: Context | None = None) -> None:  # noqa: ARG002
+        if _is_root_span(span):
+            formatted_span_id = format_span_id(span.context.span_id)
+            thread_id = threading.get_ident()
+
+            span.set_attribute(OTEL_PROFILE_ID_KEY, formatted_span_id)
+            pyroscope.add_thread_tag(thread_id, PYROSCOPE_SPAN_ID_KEY, formatted_span_id)
+            pyroscope.add_thread_tag(thread_id, PYROSCOPE_SPAN_NAME_KEY, span.name)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        if _is_root_span(span):
+            thread_id = threading.get_ident()
+            pyroscope.remove_thread_tag(thread_id, PYROSCOPE_SPAN_ID_KEY, format_span_id(span.context.span_id))
+            pyroscope.remove_thread_tag(thread_id, PYROSCOPE_SPAN_NAME_KEY, span.name)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:  # noqa: ARG002  # pragma: no cover
+        return True

pyproject.toml

@@ -42,7 +42,6 @@ dependencies = [
     "sentry-sdk>=2.7",
     "structlog>=24",
     "pyroscope-io; platform_system != 'Windows'",
-    "pyroscope-otel; platform_system != 'Windows'",
 ]
 dynamic = ["version"]
 authors = [{ name = "community-of-python" }]

tests/instruments/test_pyroscope.py

@@ -1,12 +1,19 @@
+import typing
+from unittest import mock
+from unittest.mock import Mock
+
+import fastapi
 import pydantic
 import pytest
+from fastapi.testclient import TestClient as FastAPITestClient
 
-from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpentelemetryInstrument
+from microbootstrap.bootstrappers.fastapi import FastApiOpentelemetryInstrument
+from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig
 from microbootstrap.instruments.pyroscope_instrument import PyroscopeConfig, PyroscopeInstrument
 
 
 try:
-    from pyroscope.otel import PyroscopeSpanProcessor  # type: ignore[import-untyped]
+    import pyroscope  # type: ignore[import-untyped]  # noqa: F401
 except ImportError:  # pragma: no cover
     pytest.skip("pyroscope is not installed", allow_module_level=True)
 
@@ -26,12 +33,26 @@ def test_not_ready(self) -> None:
         instrument = PyroscopeInstrument(PyroscopeConfig(pyroscope_endpoint=None))
         assert not instrument.is_ready()
 
-    def test_opentelemetry_includes_pyroscope(self) -> None:
-        otel_instrument = OpentelemetryInstrument(
-            OpentelemetryConfig(pyroscope_endpoint=pydantic.HttpUrl("http://localhost:4040"))
+    def test_opentelemetry_includes_pyroscope_2(
+        self, monkeypatch: pytest.MonkeyPatch, minimal_opentelemetry_config: OpentelemetryConfig
+    ) -> None:
+        monkeypatch.setattr("opentelemetry.sdk.trace.TracerProvider.shutdown", Mock())
+        monkeypatch.setattr("pyroscope.add_thread_tag", add_thread_tag_mock := Mock())
+        monkeypatch.setattr("pyroscope.remove_thread_tag", remove_thread_tag_mock := Mock())
+
+        minimal_opentelemetry_config.pyroscope_endpoint = pydantic.HttpUrl("http://localhost:4040")
+
+        opentelemetry_instrument: typing.Final = FastApiOpentelemetryInstrument(minimal_opentelemetry_config)
+        opentelemetry_instrument.bootstrap()
+        fastapi_application: typing.Final = opentelemetry_instrument.bootstrap_after(fastapi.FastAPI())
+
+        @fastapi_application.get("/test-handler")
+        async def test_handler() -> None: ...
+
+        FastAPITestClient(app=fastapi_application).get("/test-handler")
+
+        assert (
+            add_thread_tag_mock.mock_calls
+            == remove_thread_tag_mock.mock_calls
+            == [mock.call(mock.ANY, "span_id", mock.ANY), mock.call(mock.ANY, "span_name", "GET /test-handler")]
         )
-        otel_instrument.bootstrap()
-        assert PyroscopeSpanProcessor in {
-            type(one_span_processor)
-            for one_span_processor in otel_instrument.tracer_provider._active_span_processor._span_processors  # noqa: SLF001
-        }