diff --git a/microbootstrap/instruments/opentelemetry_instrument.py b/microbootstrap/instruments/opentelemetry_instrument.py index 6a91309..cb96ced 100644 --- a/microbootstrap/instruments/opentelemetry_instrument.py +++ b/microbootstrap/instruments/opentelemetry_instrument.py @@ -2,7 +2,6 @@ import dataclasses import logging import os -import threading import typing import pydantic @@ -200,17 +199,14 @@ 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: typing.Final = format_span_id(span.context.span_id) - thread_id: typing.Final = 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) + pyroscope.add_thread_tag(PYROSCOPE_SPAN_ID_KEY, formatted_span_id) + pyroscope.add_thread_tag(PYROSCOPE_SPAN_NAME_KEY, span.name) def on_end(self, span: ReadableSpan) -> None: if _is_root_span(span): - thread_id: typing.Final = 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) + 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: # noqa: ARG002 # pragma: no cover return True diff --git a/tests/instruments/test_pyroscope.py b/tests/instruments/test_pyroscope.py index 04c5da9..0c31b27 100644 --- a/tests/instruments/test_pyroscope.py +++ b/tests/instruments/test_pyroscope.py @@ -13,7 +13,7 @@ try: - import pyroscope # type: ignore[import-untyped] # noqa: F401 + import pyroscope # type: ignore[import-untyped] except ImportError: # pragma: no cover pytest.skip("pyroscope is not installed", allow_module_level=True) @@ -33,12 +33,16 @@ def test_not_ready(self) -> None: instrument = PyroscopeInstrument(PyroscopeConfig(pyroscope_endpoint=None)) assert not instrument.is_ready() - def test_opentelemetry_includes_pyroscope_2( + def test_opentelemetry_includes_pyroscope( 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()) + monkeypatch.setattr( + "pyroscope.add_thread_tag", add_thread_tag_mock := Mock(side_effect=pyroscope.add_thread_tag) + ) + monkeypatch.setattr( + "pyroscope.remove_thread_tag", remove_thread_tag_mock := Mock(side_effect=pyroscope.remove_thread_tag) + ) minimal_opentelemetry_config.pyroscope_endpoint = pydantic.HttpUrl("http://localhost:4040") @@ -53,5 +57,5 @@ async def test_handler() -> None: ... 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")] + == [mock.call("span_id", mock.ANY), mock.call("span_name", "GET /test-handler")] )