Register atexit callback once per configure to prevent slowdown with each span (#378)

Author: alexmojaki

logfire/_internal/config.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import atexit
 import dataclasses
 import functools
 import json
@@ -13,9 +14,10 @@
 from functools import cached_property
 from pathlib import Path
 from threading import RLock, Thread
-from typing import Any, Callable, Literal, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, cast
 from urllib.parse import urljoin
 from uuid import uuid4
+from weakref import WeakSet
 
 import requests
 from opentelemetry import metrics, trace
@@ -81,6 +83,12 @@
 from .tracer import PendingSpanProcessor, ProxyTracerProvider
 from .utils import UnexpectedResponse, ensure_data_dir_exists, get_version, read_toml_file, suppress_instrumentation
 
+if TYPE_CHECKING:
+    from .main import FastLogfireSpan, LogfireSpan
+
+# NOTE: this WeakSet is the reason that FastLogfireSpan.__slots__ has a __weakref__ slot.
+OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan] = WeakSet()
+
 CREDENTIALS_FILENAME = 'logfire_credentials.json'
 """Default base URL for the Logfire API."""
 COMMON_REQUEST_HEADERS = {'User-Agent': f'logfire/{VERSION}'}
@@ -767,6 +775,15 @@ def check_token():
                 trace.set_tracer_provider(self._tracer_provider)
                 metrics.set_meter_provider(self._meter_provider)
 
+            @atexit.register
+            def _exit_open_spans():  # type: ignore[reportUnusedFunction]  # pragma: no cover
+                # Ensure that all open spans are closed when the program exits.
+                # OTEL registers its own atexit callback in the tracer/meter providers to shut them down.
+                # Registering this callback here after the OTEL one means that this runs first.
+                # Otherwise OTEL would log an error "Already shutdown, dropping span."
+                for span in list(OPEN_SPANS):
+                    span.__exit__(None, None, None)
+
             self._initialized = True
 
             # set up context propagation for ThreadPoolExecutor and ProcessPoolExecutor

logfire/_internal/main.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-import atexit
 import inspect
 import sys
 import traceback
 import warnings
-from functools import cached_property, partial
+from functools import cached_property
 from time import time
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Literal, Sequence, TypeVar, Union, cast
 
@@ -21,7 +20,7 @@
 from ..version import VERSION
 from . import async_
 from .auto_trace import AutoTraceModule, install_auto_tracing
-from .config import GLOBAL_CONFIG, LogfireConfig
+from .config import GLOBAL_CONFIG, OPEN_SPANS, LogfireConfig
 from .constants import (
     ATTRIBUTES_JSON_SCHEMA_KEY,
     ATTRIBUTES_LOG_LEVEL_NUM_KEY,
@@ -1577,20 +1576,20 @@ def shutdown(self, timeout_millis: int = 30_000, flush: bool = True) -> bool:  #
 class FastLogfireSpan:
     """A simple version of `LogfireSpan` optimized for auto-tracing."""
 
-    __slots__ = ('_span', '_token', '_atexit')
+    # __weakref__ is needed for the OPEN_SPANS WeakSet.
+    __slots__ = ('_span', '_token', '__weakref__')
 
     def __init__(self, span: trace_api.Span) -> None:
         self._span = span
         self._token = context_api.attach(trace_api.set_span_in_context(self._span))
-        self._atexit = partial(self.__exit__, None, None, None)
-        atexit.register(self._atexit)
+        OPEN_SPANS.add(self)
 
     def __enter__(self) -> FastLogfireSpan:
         return self
 
     @handle_internal_errors()
     def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
-        atexit.unregister(self._atexit)
+        OPEN_SPANS.remove(self)
         context_api.detach(self._token)
         _exit_span(self._span, exc_value)
         self._span.end()
@@ -1615,7 +1614,6 @@ def __init__(
         self._token: None | object = None
         self._span: None | trace_api.Span = None
         self.end_on_exit = True
-        self._atexit: Callable[[], None] | None = None
 
     if not TYPE_CHECKING:  # pragma: no branch
 
@@ -1633,8 +1631,7 @@ def __enter__(self) -> LogfireSpan:
             if self._token is None:  # pragma: no branch
                 self._token = context_api.attach(trace_api.set_span_in_context(self._span))
 
-            self._atexit = partial(self.__exit__, None, None, None)
-            atexit.register(self._atexit)
+            OPEN_SPANS.add(self)
 
         return self
 
@@ -1643,8 +1640,7 @@ def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseExceptio
         if self._token is None:  # pragma: no cover
             return
 
-        if self._atexit:  # pragma: no branch
-            atexit.unregister(self._atexit)
+        OPEN_SPANS.remove(self)
 
         context_api.detach(self._token)
         self._token = None
@@ -1656,8 +1652,6 @@ def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseExceptio
         if end_on_exit_:
             self.end()
 
-        self._token = None
-
     @property
     def message_template(self) -> str | None:  # pragma: no cover
         return self._get_attribute(ATTRIBUTES_MESSAGE_TEMPLATE_KEY, None)

tests/otel_integrations/test_celery.py

@@ -1,4 +1,5 @@
 import logging
+import sys
 from typing import Generator, Iterator
 
 import pytest
@@ -13,6 +14,9 @@
 import logfire
 from logfire.testing import TestExporter
 
+# TODO find a better solution
+pytestmark = pytest.mark.skipif(sys.version_info < (3, 9), reason='Redis testcontainers has problems in 3.8')
+
 
 @pytest.fixture(scope='module', autouse=True)
 def redis_container() -> Generator[RedisContainer, None, None]: