Refactor Pyroscope tags to drop thread ID usage (#139)

Author: vrslev

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

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")]
         )