Use OPEN_SPANS and custom record_exception logic for all spans (#696

)

Co-authored-by: Alex Hall <alex.mojaki@gmail.com>

logfire-api/logfire_api/_internal/config.pyi

@@ -12,11 +12,11 @@ from .exporters.quiet_metrics import QuietMetricExporter as QuietMetricExporter
 from .exporters.remove_pending import RemovePendingSpansExporter as RemovePendingSpansExporter
 from .exporters.test import TestExporter as TestExporter
 from .integrations.executors import instrument_executors as instrument_executors
-from .main import FastLogfireSpan as FastLogfireSpan, Logfire as Logfire, LogfireSpan as LogfireSpan
+from .main import Logfire as Logfire
 from .metrics import ProxyMeterProvider as ProxyMeterProvider
 from .scrubbing import BaseScrubber as BaseScrubber, NOOP_SCRUBBER as NOOP_SCRUBBER, Scrubber as Scrubber, ScrubbingOptions as ScrubbingOptions
 from .stack_info import warn_at_user_stacklevel as warn_at_user_stacklevel
-from .tracer import PendingSpanProcessor as PendingSpanProcessor, ProxyTracerProvider as ProxyTracerProvider
+from .tracer import OPEN_SPANS as OPEN_SPANS, PendingSpanProcessor as PendingSpanProcessor, ProxyTracerProvider as ProxyTracerProvider
 from .utils import SeededRandomIdGenerator as SeededRandomIdGenerator, UnexpectedResponse as UnexpectedResponse, ensure_data_dir_exists as ensure_data_dir_exists, handle_internal_errors as handle_internal_errors, read_toml_file as read_toml_file, suppress_instrumentation as suppress_instrumentation
 from _typeshed import Incomplete
 from dataclasses import dataclass
@@ -30,9 +30,7 @@ from opentelemetry.sdk.trace.id_generator import IdGenerator
 from pathlib import Path
 from typing import Any, Callable, Literal, Sequence, TypedDict
 from typing_extensions import Self, Unpack
-from weakref import WeakSet
 
-OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan]
 CREDENTIALS_FILENAME: str
 COMMON_REQUEST_HEADERS: Incomplete
 PROJECT_NAME_PATTERN: str

logfire-api/logfire_api/_internal/main.pyi

@@ -6,9 +6,9 @@ import requests
 from . import async_ as async_
 from ..version import VERSION as VERSION
 from .auto_trace import AutoTraceModule as AutoTraceModule, install_auto_tracing as install_auto_tracing
-from .config import GLOBAL_CONFIG as GLOBAL_CONFIG, LogfireConfig as LogfireConfig, OPEN_SPANS as OPEN_SPANS
+from .config import GLOBAL_CONFIG as GLOBAL_CONFIG, LogfireConfig as LogfireConfig
 from .config_params import PydanticPluginRecordValues as PydanticPluginRecordValues
-from .constants import ATTRIBUTES_JSON_SCHEMA_KEY as ATTRIBUTES_JSON_SCHEMA_KEY, ATTRIBUTES_LOG_LEVEL_NUM_KEY as ATTRIBUTES_LOG_LEVEL_NUM_KEY, ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY as ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY as ATTRIBUTES_TAGS_KEY, ATTRIBUTES_VALIDATION_ERROR_KEY as ATTRIBUTES_VALIDATION_ERROR_KEY, DISABLE_CONSOLE_KEY as DISABLE_CONSOLE_KEY, LEVEL_NUMBERS as LEVEL_NUMBERS, LevelName as LevelName, NULL_ARGS_KEY as NULL_ARGS_KEY, OTLP_MAX_INT_SIZE as OTLP_MAX_INT_SIZE, log_level_attributes as log_level_attributes
+from .constants import ATTRIBUTES_JSON_SCHEMA_KEY as ATTRIBUTES_JSON_SCHEMA_KEY, ATTRIBUTES_LOG_LEVEL_NUM_KEY as ATTRIBUTES_LOG_LEVEL_NUM_KEY, ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY as ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY as ATTRIBUTES_TAGS_KEY, DISABLE_CONSOLE_KEY as DISABLE_CONSOLE_KEY, LEVEL_NUMBERS as LEVEL_NUMBERS, LevelName as LevelName, NULL_ARGS_KEY as NULL_ARGS_KEY, OTLP_MAX_INT_SIZE as OTLP_MAX_INT_SIZE, log_level_attributes as log_level_attributes
 from .formatter import logfire_format as logfire_format, logfire_format_with_magic as logfire_format_with_magic
 from .instrument import instrument as instrument
 from .integrations.asgi import ASGIApp as ASGIApp, ASGIInstrumentKwargs as ASGIInstrumentKwargs
@@ -29,7 +29,7 @@ from .json_encoder import logfire_json_dumps as logfire_json_dumps
 from .json_schema import JsonSchemaProperties as JsonSchemaProperties, attributes_json_schema as attributes_json_schema, attributes_json_schema_properties as attributes_json_schema_properties, create_json_schema as create_json_schema
 from .metrics import ProxyMeterProvider as ProxyMeterProvider
 from .stack_info import get_user_stack_info as get_user_stack_info
-from .tracer import ProxyTracerProvider as ProxyTracerProvider
+from .tracer import ProxyTracerProvider as ProxyTracerProvider, record_exception as record_exception, set_exception_status as set_exception_status
 from .utils import SysExcInfo as SysExcInfo, get_version as get_version, handle_internal_errors as handle_internal_errors, log_internal_error as log_internal_error, uniquify_sequence as uniquify_sequence
 from django.http import HttpRequest as HttpRequest, HttpResponse as HttpResponse
 from fastapi import FastAPI
@@ -1034,7 +1034,7 @@ class LogfireSpan(ReadableSpan):
     def message(self) -> str: ...
     @message.setter
     def message(self, message: str): ...
-    def end(self) -> None:
+    def end(self, end_time: int | None = None) -> None:
         """Sets the current time as the span's end time.
 
         The span's end time is the wall time at which the operation finished.
@@ -1093,7 +1093,7 @@ class NoopSpan:
     def is_recording(self) -> bool: ...
 AttributesValueType = TypeVar('AttributesValueType', bound=Any | otel_types.AttributeValue)
 
-def user_attributes(attributes: dict[str, Any]) -> dict[str, otel_types.AttributeValue]:
+def prepare_otlp_attributes(attributes: dict[str, Any]) -> dict[str, otel_types.AttributeValue]:
     """Prepare attributes for sending to OpenTelemetry.
 
     This will convert any non-OpenTelemetry compatible types to JSON.

logfire-api/logfire_api/_internal/tracer.pyi

@@ -1,6 +1,7 @@
 import opentelemetry.trace as trace_api
 from .config import LogfireConfig as LogfireConfig
-from .constants import ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY as ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, PENDING_SPAN_NAME_SUFFIX as PENDING_SPAN_NAME_SUFFIX
+from .constants import ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY as ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_VALIDATION_ERROR_KEY as ATTRIBUTES_VALIDATION_ERROR_KEY, PENDING_SPAN_NAME_SUFFIX as PENDING_SPAN_NAME_SUFFIX, log_level_attributes as log_level_attributes
+from .utils import handle_internal_errors as handle_internal_errors
 from _typeshed import Incomplete
 from dataclasses import dataclass
 from opentelemetry import context as context_api
@@ -13,7 +14,9 @@ from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util import types as otel_types
 from threading import Lock
 from typing import Any, Callable, Mapping, Sequence
-from weakref import WeakKeyDictionary
+from weakref import WeakKeyDictionary, WeakSet
+
+OPEN_SPANS: WeakSet[_LogfireWrappedSpan]
 
 @dataclass
 class ProxyTracerProvider(TracerProvider):
@@ -32,11 +35,12 @@ class ProxyTracerProvider(TracerProvider):
     def resource(self) -> Resource: ...
     def force_flush(self, timeout_millis: int = 30000) -> bool: ...
 
-@dataclass
-class _MaybeDeterministicTimestampSpan(trace_api.Span, ReadableSpan):
+@dataclass(eq=False)
+class _LogfireWrappedSpan(trace_api.Span, ReadableSpan):
     """Span that overrides end() to use a timestamp generator if one was provided."""
     span: Span
     ns_timestamp_generator: Callable[[], int]
+    def __post_init__(self) -> None: ...
     def end(self, end_time: int | None = None) -> None: ...
     def get_span_context(self) -> SpanContext: ...
     def set_attributes(self, attributes: dict[str, otel_types.AttributeValue]) -> None: ...
@@ -85,3 +89,6 @@ def should_sample(span_context: SpanContext, attributes: Mapping[str, otel_types
     This is used to sample spans that are not sampled by the OTEL sampler.
     """
 def get_sample_rate_from_attributes(attributes: otel_types.Attributes) -> float | None: ...
+def record_exception(span: trace_api.Span, exception: BaseException, *, attributes: otel_types.Attributes = None, timestamp: int | None = None, escaped: bool = False) -> None:
+    """Similar to the OTEL SDK Span.record_exception method, with our own additions."""
+def set_exception_status(span: trace_api.Span, exception: BaseException): ...

logfire/_internal/backfill.py

@@ -24,7 +24,7 @@
 )
 from .exporters.file import FileSpanExporter
 from .formatter import logfire_format
-from .main import user_attributes
+from .main import prepare_otlp_attributes
 from .scrubbing import Scrubber
 
 try:
@@ -150,7 +150,7 @@ def write(self, data: Union[Log, Span]) -> None:
                 )
             else:
                 parent_context = None  # pragma: no cover
-            otlp_attributes = user_attributes(data.attributes)
+            otlp_attributes = prepare_otlp_attributes(data.attributes)
 
             if data.formatted_msg is None:  # pragma: no cover
                 formatted_message = logfire_format(data.msg_template, data.attributes, self.scrubber)
@@ -196,7 +196,7 @@ def write(self, data: Union[Log, Span]) -> None:
             start_timestamp = data.start_timestamp
             if start_timestamp.tzinfo is None:  # pragma: no branch
                 start_timestamp = start_timestamp.replace(tzinfo=timezone.utc)
-            otlp_attributes = user_attributes(data.log_attributes)
+            otlp_attributes = prepare_otlp_attributes(data.log_attributes)
             if data.formatted_msg is None:  # pragma: no branch
                 formatted_message = logfire_format(data.msg_template, data.log_attributes, self.scrubber)
             else:  # pragma: no cover

logfire/_internal/config.py

@@ -16,7 +16,6 @@
 from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, TypedDict, cast
 from urllib.parse import urljoin
 from uuid import uuid4
-from weakref import WeakSet
 
 import requests
 from opentelemetry import trace
@@ -87,7 +86,7 @@
 from .metrics import ProxyMeterProvider
 from .scrubbing import NOOP_SCRUBBER, BaseScrubber, Scrubber, ScrubbingOptions
 from .stack_info import warn_at_user_stacklevel
-from .tracer import PendingSpanProcessor, ProxyTracerProvider
+from .tracer import OPEN_SPANS, PendingSpanProcessor, ProxyTracerProvider
 from .utils import (
     SeededRandomIdGenerator,
     UnexpectedResponse,
@@ -98,10 +97,8 @@
 )
 
 if TYPE_CHECKING:
-    from .main import FastLogfireSpan, Logfire, LogfireSpan
+    from .main import Logfire
 
-# 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."""
@@ -948,7 +945,7 @@ def _exit_open_spans():  # type: ignore[reportUnusedFunction]  # pragma: no cove
                 # 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)
+                    span.end()
 
             self._initialized = True