diff --git a/lite_bootstrap/instruments/opentelemetry_instrument.py b/lite_bootstrap/instruments/opentelemetry_instrument.py index 9c3ef30..bbd9fa1 100644 --- a/lite_bootstrap/instruments/opentelemetry_instrument.py +++ b/lite_bootstrap/instruments/opentelemetry_instrument.py @@ -78,6 +78,9 @@ class OpenTelemetryInstrument(BaseInstrument): bootstrap_config: OpentelemetryConfig not_ready_message = "opentelemetry_endpoint is empty and opentelemetry_log_traces is False" missing_dependency_message = "opentelemetry is not installed" + _tracer_provider: "TracerProvider | None" = dataclasses.field( + default_factory=lambda: None, init=False, repr=False, compare=False + ) def is_ready(self) -> bool: return ( @@ -105,6 +108,7 @@ def bootstrap(self) -> None: ) tracer_provider = TracerProvider(resource=resource) set_tracer_provider(tracer_provider) + object.__setattr__(self, "_tracer_provider", tracer_provider) 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: @@ -133,3 +137,6 @@ def teardown(self) -> None: one_instrumentor.instrumentor.uninstrument(**one_instrumentor.additional_params) else: one_instrumentor.uninstrument() + if self._tracer_provider is not None: + self._tracer_provider.shutdown() + object.__setattr__(self, "_tracer_provider", None) diff --git a/tests/instruments/test_opentelemetry_instrument.py b/tests/instruments/test_opentelemetry_instrument.py index 0fbccd7..3f3e68b 100644 --- a/tests/instruments/test_opentelemetry_instrument.py +++ b/tests/instruments/test_opentelemetry_instrument.py @@ -1,3 +1,5 @@ +from unittest.mock import patch + from lite_bootstrap.instruments.opentelemetry_instrument import ( InstrumentorWithParams, OpentelemetryConfig, @@ -32,3 +34,18 @@ def test_opentelemetry_instrument_empty_instruments() -> None: opentelemetry_instrument.bootstrap() finally: opentelemetry_instrument.teardown() + + +def test_opentelemetry_instrument_teardown_shuts_down_tracer_provider() -> None: + instrument = OpenTelemetryInstrument( + bootstrap_config=OpentelemetryConfig(opentelemetry_log_traces=True), + ) + instrument.bootstrap() + tracer_provider = instrument._tracer_provider # noqa: SLF001 + assert tracer_provider is not None + + with patch.object(tracer_provider, "shutdown") as mock_shutdown: + instrument.teardown() + + mock_shutdown.assert_called_once_with() + assert instrument._tracer_provider is None # noqa: SLF001