From 7bda67a3b812477faaa6a6ce5c26b42990d1fe9d Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Fri, 28 Aug 2026 17:49:22 +0900 Subject: [PATCH] fix(tracing): close and recover the default exporter safely Co-authored-by: abhay-codes07 Co-authored-by: rajarshidattapy --- src/agents/tracing/__init__.py | 8 +- src/agents/tracing/processors.py | 302 +++++++++++++++----- src/agents/tracing/provider.py | 18 ++ src/agents/tracing/setup.py | 149 +++++++++- tests/test_trace_processor.py | 474 ++++++++++++++++++++++++++++++- tests/tracing/test_setup.py | 8 + 6 files changed, 878 insertions(+), 81 deletions(-) diff --git a/src/agents/tracing/__init__.py b/src/agents/tracing/__init__.py index 28b2f28bc8..208b5c4a65 100644 --- a/src/agents/tracing/__init__.py +++ b/src/agents/tracing/__init__.py @@ -19,9 +19,9 @@ turn_span, ) from .processor_interface import TracingProcessor -from .processors import default_exporter +from .processors import _set_default_exporter_api_key, default_exporter as default_exporter from .provider import TraceProvider -from .setup import get_trace_provider, set_trace_provider +from .setup import get_trace_provider, replace_trace_processors, set_trace_provider from .span_data import ( AgentSpanData, CustomSpanData, @@ -102,7 +102,7 @@ def set_trace_processors(processors: list[TracingProcessor]) -> None: """ Set the list of trace processors. This will replace the current list of processors. """ - get_trace_provider().set_processors(processors) + replace_trace_processors(processors) def set_tracing_disabled(disabled: bool) -> None: @@ -116,7 +116,7 @@ def set_tracing_export_api_key(api_key: str) -> None: """ Set the OpenAI API key for the backend exporter. """ - default_exporter().set_api_key(api_key) + _set_default_exporter_api_key(api_key) def flush_traces() -> None: diff --git a/src/agents/tracing/processors.py b/src/agents/tracing/processors.py index b61f3e7976..f849006092 100644 --- a/src/agents/tracing/processors.py +++ b/src/agents/tracing/processors.py @@ -85,6 +85,8 @@ def __init__( self.base_delay = base_delay self.max_delay = max_delay self._shutdown_event = threading.Event() + self._close_lock = threading.Lock() + self._closed = False # Keep a client open for connection pooling across multiple export calls self._client = httpx2.Client(timeout=httpx2.Timeout(timeout=60, connect=5.0)) @@ -530,8 +532,17 @@ def _sanitize_json_compatible_value(self, value: Any, seen_ids: set[int] | None return sanitized_list return self._UNSERIALIZABLE + @property + def is_shut_down(self) -> bool: + """Whether shutdown or close made this exporter terminal.""" + return self._shutdown_event.is_set() or self._closed + def close(self): - """Close the underlying HTTP client.""" + """Close the underlying HTTP client exactly once.""" + with self._close_lock: + if self._closed: + return + self._closed = True self._client.close() def _request_shutdown(self) -> None: @@ -552,6 +563,8 @@ def __init__( max_batch_size: int = 128, schedule_delay: float = 5.0, export_trigger_ratio: float = 0.7, + *, + _owns_exporter: bool = False, ): """ Args: @@ -563,6 +576,7 @@ def __init__( export_trigger_ratio: The ratio of the queue size at which we will trigger an export. """ self._exporter = exporter + self._owns_exporter = _owns_exporter self._queue: queue.Queue[Trace | Span[Any]] = queue.Queue(maxsize=max_queue_size) self._max_queue_size = max_queue_size self._max_batch_size = max_batch_size @@ -581,27 +595,27 @@ def __init__( self._export_lock = threading.Lock() self._shutdown_deadline: float | None = None - def _ensure_thread_started(self) -> None: - # Fast path without holding the lock - if self._worker_thread and self._worker_thread.is_alive(): - return + def _begin_shutdown(self) -> threading.Thread | None: + """Close admission and return the worker that owns any accepted work.""" + with self._thread_start_lock: + self._shutdown_event.set() + return self._worker_thread - # Double-checked locking to avoid starting multiple threads + def _enqueue(self, item: Trace | Span[Any], full_message: str) -> None: + """Publish work atomically with lazy worker creation and shutdown admission.""" with self._thread_start_lock: - if self._worker_thread and self._worker_thread.is_alive(): + if self._shutdown_event.is_set(): return - - self._worker_thread = threading.Thread(target=self._run, daemon=True) - self._worker_thread.start() + if not self._worker_thread or not self._worker_thread.is_alive(): + self._worker_thread = threading.Thread(target=self._run, daemon=True) + self._worker_thread.start() + try: + self._queue.put_nowait(item) + except queue.Full: + logger.warning(full_message) def on_trace_start(self, trace: Trace) -> None: - # Ensure the background worker is running before we enqueue anything. - self._ensure_thread_started() - - try: - self._queue.put_nowait(trace) - except queue.Full: - logger.warning("Queue is full, dropping trace.") + self._enqueue(trace, "Queue is full, dropping trace.") def on_trace_end(self, trace: Trace) -> None: # We send traces via on_trace_start, so we don't need to do anything here. @@ -612,19 +626,21 @@ def on_span_start(self, span: Span[Any]) -> None: pass def on_span_end(self, span: Span[Any]) -> None: - # Ensure the background worker is running before we enqueue anything. - self._ensure_thread_started() - - try: - self._queue.put_nowait(span) - except queue.Full: - logger.warning("Queue is full, dropping span.") + self._enqueue(span, "Queue is full, dropping span.") + + @property + def is_shut_down(self) -> bool: + """Whether shutdown has started. Shutdown is terminal for this processor.""" + exporter_is_shut_down = getattr(self._exporter, "is_shut_down", False) + return self._shutdown_event.is_set() or ( + self._owns_exporter and exporter_is_shut_down is True + ) def shutdown(self, timeout: float | None = None): """ Called when the application stops. We signal our thread to stop, then join it. """ - self._shutdown_event.set() + worker_thread = self._begin_shutdown() if timeout is not None: request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None) if callable(request_exporter_shutdown): @@ -634,38 +650,63 @@ def shutdown(self, timeout: float | None = None): self._shutdown_deadline = deadline # Only join if we ever started the background thread; otherwise flush synchronously. - if self._worker_thread and self._worker_thread.is_alive(): - self._worker_thread.join(timeout=timeout) - if self._worker_thread.is_alive(): - logger.warning( - "[non-fatal] Tracing: shutdown timeout reached; dropping queued traces." - ) + if worker_thread is not None: + if worker_thread.is_alive(): + worker_thread.join(timeout=timeout) + if worker_thread.is_alive(): + logger.warning( + "[non-fatal] Tracing: shutdown timeout reached; dropping queued traces." + ) + return else: # No background thread: process any remaining items synchronously. self._export_batches(deadline=deadline) + self._close_owned_exporter() + def force_flush(self): """ Forces an immediate flush of all queued spans. """ - self._export_batches() + if not self._shutdown_event.is_set(): + self._export_batches() def _run(self): - while not self._shutdown_event.is_set(): - current_time = time.monotonic() - queue_size = self._queue.qsize() - - # If it's time for a scheduled flush or queue is above the trigger threshold - if current_time >= self._next_export_time or queue_size >= self._export_trigger_size: - self._export_batches() - # Reset the next scheduled flush time - self._next_export_time = time.monotonic() + self._schedule_delay - else: - # Sleep a short interval so we don't busy-wait. - time.sleep(0.2) + try: + while not self._shutdown_event.is_set(): + current_time = time.monotonic() + queue_size = self._queue.qsize() + + # If it's time for a scheduled flush or queue is above the trigger threshold + if ( + current_time >= self._next_export_time + or queue_size >= self._export_trigger_size + ): + self._export_batches() + # Reset the next scheduled flush time + self._next_export_time = time.monotonic() + self._schedule_delay + else: + # Sleep a short interval so we don't busy-wait. + time.sleep(0.2) - # Final drain after shutdown - self._export_batches(deadline=self._shutdown_deadline) + # Final drain after shutdown + self._export_batches(deadline=self._shutdown_deadline) + finally: + self._close_owned_exporter() + + def _close_owned_exporter(self) -> None: + """Close only the exporter owned by the module-created default processor.""" + if not self._owns_exporter: + return + close = getattr(self._exporter, "close", None) + if not callable(close): + return + try: + close() + except Exception as exc: + log_model_and_tool_action_error( + logger, "[non-fatal] Tracing: error closing exporter on shutdown", exc + ) def _export_batches(self, deadline: float | None = None): """Drains the queue and exports in batches of up to `max_batch_size` until the queue @@ -722,42 +763,171 @@ def _export_batches(self, deadline: float | None = None): _global_exporter: BackendSpanExporter | None = None _global_processor: BatchTraceProcessor | None = None _global_lock = threading.Lock() +_DEFAULT_STACK_ATEXIT_SHUTDOWN = False +_global_exporter_attached = False +_retirement_lock = threading.Lock() +_retiring_default_processors: dict[BatchTraceProcessor, threading.Thread] = {} -def default_exporter() -> BackendSpanExporter: - """The default exporter, which exports traces and spans to the backend in batches.""" +def _begin_default_stack_atexit_shutdown() -> None: + """Prevent a later atexit callback from rebuilding terminal defaults.""" + global _DEFAULT_STACK_ATEXIT_SHUTDOWN + + with _global_lock: + _DEFAULT_STACK_ATEXIT_SHUTDOWN = True + + +def _close_unattached_default_exporter() -> None: + """Close a configured default exporter that never gained an owning processor.""" + with _global_lock: + if _global_processor is not None or _global_exporter_attached: + return + exporter = _global_exporter + + if exporter is not None: + try: + exporter.close() + except Exception as exc: + log_model_and_tool_action_error( + logger, "[non-fatal] Tracing: error closing unattached exporter", exc + ) + + +def _wait_for_retiring_default_processors(timeout: float | None) -> None: + """Give detached SDK defaults the same bounded atexit drain window.""" + deadline = None if timeout is None else time.monotonic() + timeout + with _retirement_lock: + threads = tuple(_retiring_default_processors.values()) + + for thread in threads: + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + if remaining is not None and remaining <= 0: + return + thread.join(timeout=remaining) + + +def _replacement_exporter(exporter: BackendSpanExporter) -> BackendSpanExporter: + """Create a live default exporter with the prior default exporter's configuration.""" + return BackendSpanExporter( + api_key=exporter._api_key, + organization=exporter._organization, + project=exporter._project, + endpoint=exporter.endpoint, + max_retries=exporter.max_retries, + base_delay=exporter.base_delay, + max_delay=exporter.max_delay, + ) + + +def _replace_shut_down_defaults() -> None: + """Replace terminal cached defaults while retaining their configured values. + + Callers must hold the default-stack lock. + """ global _global_exporter + global _global_processor + global _global_exporter_attached + processor = _global_processor exporter = _global_exporter - if exporter is not None: - return exporter + if processor is not None and processor.is_shut_down: + if not processor._shutdown_event.is_set(): + processor.shutdown(timeout=0.0) + _global_processor = None + if exporter is not None: + _global_exporter = _replacement_exporter(exporter) + _global_exporter_attached = False + return + if exporter is not None and exporter.is_shut_down: + if processor is not None: + processor.shutdown(timeout=0.0) + _global_exporter = _replacement_exporter(exporter) + _global_processor = None + _global_exporter_attached = False + + +def default_exporter() -> BackendSpanExporter: + """The default exporter, which exports traces and spans to the backend in batches.""" + global _global_exporter with _global_lock: - exporter = _global_exporter - if exporter is None: - exporter = BackendSpanExporter() - _global_exporter = exporter + if _DEFAULT_STACK_ATEXIT_SHUTDOWN and _global_exporter is None: + raise RuntimeError("Tracing shutdown has started; default exporter is unavailable.") + if not _DEFAULT_STACK_ATEXIT_SHUTDOWN: + _replace_shut_down_defaults() + if _global_exporter is None: + _global_exporter = BackendSpanExporter() + return _global_exporter - return exporter + +def _set_default_exporter_api_key(api_key: str) -> None: + """Update the active default exporter atomically with terminal replacement.""" + global _global_exporter + + with _global_lock: + if _DEFAULT_STACK_ATEXIT_SHUTDOWN: + return + _replace_shut_down_defaults() + if _global_exporter is None: + _global_exporter = BackendSpanExporter() + _global_exporter.set_api_key(api_key) def default_processor() -> BatchTraceProcessor: """The default processor, which exports traces and spans to the backend in batches.""" global _global_exporter global _global_processor + global _global_exporter_attached - processor = _global_processor - if processor is not None: - return processor + with _global_lock: + if _DEFAULT_STACK_ATEXIT_SHUTDOWN and _global_processor is None: + raise RuntimeError("Tracing shutdown has started; default processor is unavailable.") + if not _DEFAULT_STACK_ATEXIT_SHUTDOWN: + _replace_shut_down_defaults() + if _global_processor is None: + if _global_exporter is None: + _global_exporter = BackendSpanExporter() + _global_processor = BatchTraceProcessor(_global_exporter, _owns_exporter=True) + _global_exporter_attached = True + return _global_processor + + +def _detach_owned_default_processor(processor: TracingProcessor) -> None: + """Publish a fresh cache and close admission on the retired default.""" + global _global_exporter + global _global_processor + global _global_exporter_attached + + if not isinstance(processor, BatchTraceProcessor) or not processor._owns_exporter: + return with _global_lock: - processor = _global_processor - if processor is None: + if _global_processor is processor: exporter = _global_exporter - if exporter is None: - exporter = BackendSpanExporter() - _global_exporter = exporter - processor = BatchTraceProcessor(exporter) - _global_processor = processor + _global_processor = None + if exporter is processor._exporter: + _global_exporter = _replacement_exporter(exporter) + _global_exporter_attached = False + processor._begin_shutdown() + - return processor +def _retire_owned_default_processor(processor: TracingProcessor) -> None: + """Drain an SDK-owned default asynchronously without delaying reconfiguration.""" + if not isinstance(processor, BatchTraceProcessor) or not processor._owns_exporter: + return + + _detach_owned_default_processor(processor) + + def retire() -> None: + try: + processor.shutdown(timeout=None) + finally: + with _retirement_lock: + _retiring_default_processors.pop(processor, None) + + thread = threading.Thread(target=retire, daemon=True) + with _retirement_lock: + if processor in _retiring_default_processors: + return + _retiring_default_processors[processor] = thread + thread.start() diff --git a/src/agents/tracing/provider.py b/src/agents/tracing/provider.py index b0e10b0bd2..10516fe999 100644 --- a/src/agents/tracing/provider.py +++ b/src/agents/tracing/provider.py @@ -114,6 +114,13 @@ def set_processors(self, processors: list[TracingProcessor]): with self._lock: self._processors = tuple(processors) + def _replace_processor(self, old: TracingProcessor, new: TracingProcessor) -> None: + """Replace one registered processor without disturbing registration order.""" + with self._lock: + self._processors = tuple( + new if processor is old else processor for processor in self._processors + ) + def on_trace_start(self, trace: Trace) -> None: """ Called when a trace is started. @@ -315,8 +322,19 @@ def set_processors(self, processors: list[TracingProcessor]): """ Set the list of processors. This will replace the current list of processors. """ + from .setup import replace_trace_processors_for_provider + + if not replace_trace_processors_for_provider(self, processors): + self._set_processors(processors) + + def _set_processors(self, processors: list[TracingProcessor]) -> None: + """Replace processors after setup has handled SDK-default ownership.""" self._multi_processor.set_processors(processors) + def _replace_processor(self, old: TracingProcessor, new: TracingProcessor) -> None: + """Replace one registered processor without rebuilding this provider.""" + self._multi_processor._replace_processor(old, new) + def get_current_trace(self) -> Trace | None: """ Returns the currently active trace, if any. diff --git a/src/agents/tracing/setup.py b/src/agents/tracing/setup.py index 0ec72de239..03c3cee0a3 100644 --- a/src/agents/tracing/setup.py +++ b/src/agents/tracing/setup.py @@ -5,36 +5,101 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from .processor_interface import TracingProcessor from .provider import TraceProvider _DEFAULT_SHUTDOWN_TIMEOUT = 5.0 GLOBAL_TRACE_PROVIDER: TraceProvider | None = None _GLOBAL_TRACE_PROVIDER_LOCK = threading.Lock() _SHUTDOWN_HANDLER_REGISTERED = False +_DEFAULT_PROCESSOR: TracingProcessor | None = None +_SDK_DEFAULT_PROVIDER: TraceProvider | None = None +_SDK_DEFAULT_PROCESSOR: TracingProcessor | None = None +_ATEXIT_SHUTDOWN_STARTED = False + + +def _default_processor_is_terminal() -> bool: + """Return whether the SDK-owned processor needs supported-use recovery.""" + if _ATEXIT_SHUTDOWN_STARTED: + return False + return getattr(_DEFAULT_PROCESSOR, "is_shut_down", False) is True + + +def _recover_default_processor(provider: TraceProvider) -> None: + """Replace only the SDK-owned terminal processor in the existing provider.""" + global _DEFAULT_PROCESSOR + global _SDK_DEFAULT_PROCESSOR + + from .processors import default_processor + from .provider import DefaultTraceProvider + + stale = _DEFAULT_PROCESSOR + fresh = default_processor() + if fresh is stale: + return + if stale is not None and isinstance(provider, DefaultTraceProvider): + provider._replace_processor(stale, fresh) + _DEFAULT_PROCESSOR = fresh + _SDK_DEFAULT_PROCESSOR = fresh def _shutdown_global_trace_provider() -> None: - provider = GLOBAL_TRACE_PROVIDER - if provider is not None: - from .provider import DefaultTraceProvider + global _ATEXIT_SHUTDOWN_STARTED + + from .processors import _begin_default_stack_atexit_shutdown + + with _GLOBAL_TRACE_PROVIDER_LOCK: + _ATEXIT_SHUTDOWN_STARTED = True + _begin_default_stack_atexit_shutdown() + provider = GLOBAL_TRACE_PROVIDER + try: + if provider is not None: + from .provider import DefaultTraceProvider + + if isinstance(provider, DefaultTraceProvider): + provider.shutdown(timeout=_DEFAULT_SHUTDOWN_TIMEOUT) + else: + provider.shutdown() + finally: + from .processors import ( + _close_unattached_default_exporter, + _wait_for_retiring_default_processors, + ) - if isinstance(provider, DefaultTraceProvider): - provider.shutdown(timeout=_DEFAULT_SHUTDOWN_TIMEOUT) - return - provider.shutdown() + _wait_for_retiring_default_processors(_DEFAULT_SHUTDOWN_TIMEOUT) + _close_unattached_default_exporter() def set_trace_provider(provider: TraceProvider) -> None: """Set the global trace provider used by tracing utilities.""" global GLOBAL_TRACE_PROVIDER global _SHUTDOWN_HANDLER_REGISTERED + global _DEFAULT_PROCESSOR + global _SDK_DEFAULT_PROVIDER + global _SDK_DEFAULT_PROCESSOR + retired_default: TracingProcessor | None = None with _GLOBAL_TRACE_PROVIDER_LOCK: + if provider is not GLOBAL_TRACE_PROVIDER: + retired_default = _DEFAULT_PROCESSOR + if retired_default is not None: + from .processors import _detach_owned_default_processor + + _detach_owned_default_processor(retired_default) GLOBAL_TRACE_PROVIDER = provider + if retired_default is not None: + _DEFAULT_PROCESSOR = None + if provider is _SDK_DEFAULT_PROVIDER: + _DEFAULT_PROCESSOR = _SDK_DEFAULT_PROCESSOR if not _SHUTDOWN_HANDLER_REGISTERED: atexit.register(_shutdown_global_trace_provider) _SHUTDOWN_HANDLER_REGISTERED = True + if retired_default is not None: + from .processors import _retire_owned_default_processor + + _retire_owned_default_processor(retired_default) + def get_trace_provider() -> TraceProvider: """Get the global trace provider used by tracing utilities. @@ -44,23 +109,87 @@ def get_trace_provider() -> TraceProvider: """ global GLOBAL_TRACE_PROVIDER global _SHUTDOWN_HANDLER_REGISTERED + global _DEFAULT_PROCESSOR + global _SDK_DEFAULT_PROVIDER + global _SDK_DEFAULT_PROCESSOR provider = GLOBAL_TRACE_PROVIDER - if provider is not None: + if provider is not None and not _default_processor_is_terminal(): return provider with _GLOBAL_TRACE_PROVIDER_LOCK: provider = GLOBAL_TRACE_PROVIDER if provider is None: - from .processors import default_processor from .provider import DefaultTraceProvider provider = DefaultTraceProvider() - provider.register_processor(default_processor()) GLOBAL_TRACE_PROVIDER = provider + if not _ATEXIT_SHUTDOWN_STARTED: + from .processors import default_processor + + processor = default_processor() + provider.register_processor(processor) + _DEFAULT_PROCESSOR = processor + _SDK_DEFAULT_PROVIDER = provider + _SDK_DEFAULT_PROCESSOR = processor + elif _default_processor_is_terminal(): + _recover_default_processor(provider) if not _SHUTDOWN_HANDLER_REGISTERED: atexit.register(_shutdown_global_trace_provider) _SHUTDOWN_HANDLER_REGISTERED = True return provider + + +def replace_trace_processors(processors: list[TracingProcessor]) -> None: + """Replace processors and retire a removed SDK-owned default processor.""" + get_trace_provider().set_processors(processors) + + +def _contains_processor_identity( + processors: list[TracingProcessor], target: TracingProcessor +) -> bool: + """Return whether the exact SDK-owned processor remains registered.""" + return any(processor is target for processor in processors) + + +def replace_trace_processors_for_provider( + provider: TraceProvider, processors: list[TracingProcessor] +) -> bool: + """Handle direct replacement on the currently registered SDK default provider.""" + global _DEFAULT_PROCESSOR + global _SDK_DEFAULT_PROCESSOR + + with _GLOBAL_TRACE_PROVIDER_LOCK: + from .provider import DefaultTraceProvider + + if not isinstance(provider, DefaultTraceProvider): + return False + if provider is not GLOBAL_TRACE_PROVIDER: + if ( + provider is _SDK_DEFAULT_PROVIDER + and _SDK_DEFAULT_PROCESSOR is not None + and not _contains_processor_identity(processors, _SDK_DEFAULT_PROCESSOR) + ): + _SDK_DEFAULT_PROCESSOR = None + return False + if _DEFAULT_PROCESSOR is None: + return False + provider._set_processors(processors) + default_processor = _DEFAULT_PROCESSOR + default_removed = default_processor is not None and not _contains_processor_identity( + processors, default_processor + ) + if default_removed: + from .processors import _detach_owned_default_processor + + _detach_owned_default_processor(default_processor) + _DEFAULT_PROCESSOR = None + _SDK_DEFAULT_PROCESSOR = None + + if default_removed: + from .processors import _retire_owned_default_processor + + _retire_owned_default_processor(default_processor) + return True diff --git a/tests/test_trace_processor.py b/tests/test_trace_processor.py index 07e975ccb9..6099c00430 100644 --- a/tests/test_trace_processor.py +++ b/tests/test_trace_processor.py @@ -12,7 +12,14 @@ import pytest import agents._debug as _debug -from agents.tracing import flush_traces, get_trace_provider +from agents.tracing import ( + flush_traces, + get_trace_provider, + processors as tracing_processors, + set_trace_processors, + set_tracing_export_api_key, + setup as tracing_setup, +) from agents.tracing.processor_interface import TracingExporter, TracingProcessor from agents.tracing.processors import BackendSpanExporter, BatchTraceProcessor, ConsoleSpanExporter from agents.tracing.provider import DefaultTraceProvider, TraceProvider @@ -612,6 +619,471 @@ def test_batch_trace_processor_shutdown_without_timeout_preserves_export_retries exporter.close() +def _reset_default_tracing_stack(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(tracing_processors, "_global_exporter", None) + monkeypatch.setattr(tracing_processors, "_global_processor", None) + monkeypatch.setattr(tracing_processors, "_global_exporter_attached", False) + monkeypatch.setattr(tracing_processors, "_DEFAULT_STACK_ATEXIT_SHUTDOWN", False) + monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", None) + monkeypatch.setattr(tracing_setup, "_DEFAULT_PROCESSOR", None) + monkeypatch.setattr(tracing_setup, "_SDK_DEFAULT_PROVIDER", None) + monkeypatch.setattr(tracing_setup, "_SDK_DEFAULT_PROCESSOR", None) + monkeypatch.setattr(tracing_setup, "_ATEXIT_SHUTDOWN_STARTED", False) + monkeypatch.setattr(tracing_setup, "_SHUTDOWN_HANDLER_REGISTERED", True) + + +def _wait_for_close(mock_client: MagicMock, timeout: float = 2.0) -> None: + deadline = time.monotonic() + timeout + while mock_client.return_value.close.call_count == 0 and time.monotonic() < deadline: + time.sleep(0.01) + assert mock_client.return_value.close.call_count > 0 + + +@patch("httpx2.Client") +def test_default_exporter_closes_once_after_no_timeout_shutdown(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + exporter = tracing_processors.default_exporter() + processor = tracing_processors.default_processor() + processor._queue.put_nowait(get_span(processor)) + + processor.shutdown(timeout=None) + processor.shutdown(timeout=None) + + mock_client.return_value.close.assert_called_once() + replacement = tracing_processors.default_exporter() + assert replacement is not exporter + replacement.close() + + +@patch("httpx2.Client") +def test_timed_default_shutdown_does_not_close_under_surviving_export(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + exporter = tracing_processors.default_exporter() + processor = tracing_processors.default_processor() + export_started = threading.Event() + release_export = threading.Event() + + def blocked_export(_items: list[Trace | Span[Any]]) -> None: + export_started.set() + assert release_export.wait(timeout=2.0) + + monkeypatch.setattr(exporter, "export", blocked_export) + processor._export_trigger_size = 1 + processor.on_span_end(get_span(processor)) + assert export_started.wait(timeout=2.0) + + processor.shutdown(timeout=0.01) + mock_client.return_value.close.assert_not_called() + + release_export.set() + assert processor._worker_thread is not None + processor._worker_thread.join(timeout=2.0) + assert not processor._worker_thread.is_alive() + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_shutdown_waits_for_callback_admission_before_closing(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + processor = tracing_processors.default_processor() + put_started = threading.Event() + release_put = threading.Event() + original_put = processor._queue.put_nowait + + def blocked_put(item: Trace | Span[Any]) -> None: + put_started.set() + assert release_put.wait(timeout=2.0) + original_put(item) + + monkeypatch.setattr(processor._queue, "put_nowait", blocked_put) + callback_thread = threading.Thread(target=processor.on_span_end, args=(get_span(processor),)) + callback_thread.start() + assert put_started.wait(timeout=2.0) + + shutdown_thread = threading.Thread(target=processor.shutdown, kwargs={"timeout": None}) + shutdown_thread.start() + time.sleep(0.05) + assert shutdown_thread.is_alive() + mock_client.return_value.close.assert_not_called() + + release_put.set() + callback_thread.join(timeout=2.0) + shutdown_thread.join(timeout=2.0) + assert not callback_thread.is_alive() + assert not shutdown_thread.is_alive() + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_shutdown_leaves_caller_injected_exporter_open(mock_client): + exporter = BackendSpanExporter(api_key="test_key") + processor = BatchTraceProcessor(exporter=exporter) + + processor.shutdown(timeout=None) + + mock_client.return_value.close.assert_not_called() + exporter.close() + + +@patch("httpx2.Client") +def test_repeated_shutdown_does_not_drain_after_owned_worker_closed(mock_client) -> None: + exporter = BackendSpanExporter(api_key="test_key") + processor = BatchTraceProcessor(exporter=exporter, schedule_delay=30.0, _owns_exporter=True) + processor.on_span_end(get_span(processor)) + + processor.shutdown(timeout=0.0) + assert processor._worker_thread is not None + processor._worker_thread.join(timeout=2.0) + processor.shutdown(timeout=None) + + mock_client.return_value.post.assert_not_called() + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_set_trace_processors_retires_removed_owned_default(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + exporter = tracing_processors.default_exporter() + provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + custom_processor = MagicMock() + custom_processor.__eq__.return_value = True + + set_trace_processors([custom_processor]) + + assert len(provider._multi_processor._processors) == 1 + assert provider._multi_processor._processors[0] is custom_processor + assert tracing_setup._DEFAULT_PROCESSOR is None + _wait_for_close(mock_client) + mock_client.return_value.close.assert_called_once() + assert tracing_processors.default_exporter() is not exporter + + +@patch("httpx2.Client") +def test_set_trace_processors_drains_queued_default_before_retiring(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + response = MagicMock() + response.status_code = 200 + mock_client.return_value.post.return_value = response + exporter = tracing_processors.default_exporter() + exporter.set_api_key("test_key") + processor = tracing_processors.default_processor() + tracing_setup.get_trace_provider() + processor._queue.put_nowait(get_span(processor)) + + set_trace_processors([MagicMock()]) + + _wait_for_close(mock_client) + mock_client.return_value.post.assert_called_once() + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_direct_provider_set_processors_retires_owned_default(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + custom_processor = MagicMock() + + provider.set_processors([custom_processor]) + + assert provider._multi_processor._processors == (custom_processor,) + assert tracing_setup._DEFAULT_PROCESSOR is None + _wait_for_close(mock_client) + mock_client.return_value.close.assert_called_once() + + +def test_direct_provider_set_processors_without_default_does_not_wait_for_callback(): + provider = DefaultTraceProvider() + callback_started = threading.Event() + release_callback = threading.Event() + blocking_processor = MagicMock() + replacement_processor = MagicMock() + + def block_callback(_span: Span[Any]) -> None: + callback_started.set() + assert release_callback.wait(timeout=2.0) + + blocking_processor.on_span_end.side_effect = block_callback + provider.set_processors([blocking_processor]) + callback_thread = threading.Thread( + target=provider._multi_processor.on_span_end, + args=(get_span(blocking_processor),), + ) + callback_thread.start() + assert callback_started.wait(timeout=2.0) + + replacement_thread = threading.Thread( + target=provider.set_processors, + args=([replacement_processor],), + ) + replacement_thread.start() + replacement_thread.join(timeout=2.0) + assert not replacement_thread.is_alive() + + release_callback.set() + callback_thread.join(timeout=2.0) + assert not callback_thread.is_alive() + assert provider._multi_processor._processors == (replacement_processor,) + + +@patch("httpx2.Client") +def test_atexit_does_not_close_detached_exporter_under_live_worker(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + retired_client = MagicMock() + replacement_client = MagicMock() + mock_client.side_effect = [retired_client, replacement_client] + exporter = tracing_processors.default_exporter() + processor = tracing_processors.default_processor() + tracing_setup.get_trace_provider() + export_started = threading.Event() + release_export = threading.Event() + + def blocked_export(_items: list[Trace | Span[Any]]) -> None: + export_started.set() + assert release_export.wait(timeout=2.0) + + monkeypatch.setattr(exporter, "export", blocked_export) + processor._export_trigger_size = 1 + processor.on_span_end(get_span(processor)) + assert export_started.wait(timeout=2.0) + + replacement_thread = threading.Thread(target=set_trace_processors, args=([MagicMock()],)) + replacement_thread.start() + replacement_thread.join(timeout=2.0) + assert not replacement_thread.is_alive() + retired_client.close.assert_not_called() + + atexit_thread = threading.Thread(target=tracing_setup._shutdown_global_trace_provider) + atexit_thread.start() + time.sleep(0.05) + assert atexit_thread.is_alive() + retired_client.close.assert_not_called() + + release_export.set() + atexit_thread.join(timeout=2.0) + assert not atexit_thread.is_alive() + assert processor._worker_thread is not None + processor._worker_thread.join(timeout=2.0) + retired_client.close.assert_called_once() + replacement_client.close.assert_called_once() + + +@patch("httpx2.Client") +def test_supported_recovery_keeps_provider_state_and_exporter_configuration( + mock_client, monkeypatch +): + _reset_default_tracing_stack(monkeypatch) + exporter = tracing_processors.default_exporter() + exporter.set_api_key("trace-only-key") + exporter._organization = "org_123" + exporter._project = "proj_123" + exporter.endpoint = "https://example.test/traces" + exporter.max_retries = 7 + exporter.base_delay = 0.25 + exporter.max_delay = 2.5 + provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + custom_processor = MagicMock() + provider.register_processor(custom_processor) + provider.set_disabled(True) + stale_processor = tracing_setup._DEFAULT_PROCESSOR + + provider.shutdown(timeout=None) + recovered = tracing_setup.get_trace_provider() + + assert recovered is provider + assert provider._manual_disabled is True + assert custom_processor in provider._multi_processor._processors + assert tracing_setup._DEFAULT_PROCESSOR is not stale_processor + replacement = tracing_processors.default_exporter() + assert replacement is not exporter + assert replacement._api_key == "trace-only-key" + assert replacement._organization == "org_123" + assert replacement._project == "proj_123" + assert replacement.endpoint == "https://example.test/traces" + assert replacement.max_retries == 7 + assert replacement.base_delay == 0.25 + assert replacement.max_delay == 2.5 + assert mock_client.return_value.close.call_count == 1 + replacement.close() + + +@patch("httpx2.Client") +def test_closed_default_exporter_recovers_registered_default_processor(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + exporter = tracing_processors.default_exporter() + stale_processor = tracing_setup._DEFAULT_PROCESSOR + + exporter.close() + replacement = tracing_processors.default_exporter() + recovered = tracing_setup.get_trace_provider() + + assert recovered is provider + assert replacement is not exporter + assert tracing_setup._DEFAULT_PROCESSOR is not stale_processor + assert provider._multi_processor._processors == (tracing_setup._DEFAULT_PROCESSOR,) + assert tracing_setup._DEFAULT_PROCESSOR is tracing_processors.default_processor() + replacement.close() + + +@patch("httpx2.Client") +def test_set_tracing_export_api_key_updates_recovered_default(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + provider.shutdown(timeout=None) + + set_tracing_export_api_key("replacement-key") + + assert tracing_processors.default_exporter()._api_key == "replacement-key" + tracing_processors.default_exporter().close() + + +@patch("httpx2.Client") +def test_set_trace_provider_retires_previous_owned_default(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + tracing_setup.get_trace_provider() + custom_provider = MagicMock() + + tracing_setup.set_trace_provider(custom_provider) + + assert tracing_setup.GLOBAL_TRACE_PROVIDER is custom_provider + assert tracing_setup._DEFAULT_PROCESSOR is None + _wait_for_close(mock_client) + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_set_trace_provider_drains_queued_default_before_retiring(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + response = MagicMock() + response.status_code = 200 + mock_client.return_value.post.return_value = response + exporter = tracing_processors.default_exporter() + exporter.set_api_key("test_key") + processor = tracing_processors.default_processor() + tracing_setup.get_trace_provider() + processor._queue.put_nowait(get_span(processor)) + + tracing_setup.set_trace_provider(MagicMock()) + + _wait_for_close(mock_client) + mock_client.return_value.post.assert_called_once() + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_reinstalling_sdk_provider_recovers_its_retired_default(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + sdk_provider = tracing_setup.get_trace_provider() + stale_processor = tracing_setup._DEFAULT_PROCESSOR + custom_provider = MagicMock() + + tracing_setup.set_trace_provider(custom_provider) + tracing_setup.set_trace_provider(sdk_provider) + recovered = tracing_setup.get_trace_provider() + + assert recovered is sdk_provider + assert tracing_setup._DEFAULT_PROCESSOR is not stale_processor + assert tracing_setup._DEFAULT_PROCESSOR in sdk_provider._multi_processor._processors + _wait_for_close(mock_client) + assert mock_client.return_value.close.call_count == 1 + tracing_processors.default_exporter().close() + + +@patch("httpx2.Client") +def test_inactive_sdk_provider_replacement_does_not_restore_retired_default( + mock_client, monkeypatch +): + _reset_default_tracing_stack(monkeypatch) + sdk_provider = cast(DefaultTraceProvider, tracing_setup.get_trace_provider()) + custom_provider = MagicMock() + replacement_processor = MagicMock() + replacement_processor.__eq__.return_value = True + + tracing_setup.set_trace_provider(custom_provider) + sdk_provider.set_processors([replacement_processor]) + tracing_setup.set_trace_provider(sdk_provider) + + assert tracing_setup.get_trace_provider() is sdk_provider + assert tracing_setup._DEFAULT_PROCESSOR is None + assert len(sdk_provider._multi_processor._processors) == 1 + assert sdk_provider._multi_processor._processors[0] is replacement_processor + _wait_for_close(mock_client) + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_atexit_shutdown_does_not_recover_default_tracing(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + provider = tracing_setup.get_trace_provider() + processor = tracing_setup._DEFAULT_PROCESSOR + exporter = tracing_processors.default_exporter() + + tracing_setup._shutdown_global_trace_provider() + retrieved = tracing_setup.get_trace_provider() + + assert retrieved is provider + assert tracing_setup._DEFAULT_PROCESSOR is processor + assert tracing_setup._ATEXIT_SHUTDOWN_STARTED is True + assert tracing_processors.default_exporter() is exporter + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_atexit_before_bootstrap_does_not_create_default_stack(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + + tracing_setup._shutdown_global_trace_provider() + provider = tracing_setup.get_trace_provider() + + assert isinstance(provider, DefaultTraceProvider) + assert tracing_setup._DEFAULT_PROCESSOR is None + assert tracing_processors._global_processor is None + assert tracing_processors._global_exporter is None + mock_client.assert_not_called() + + +@patch("httpx2.Client") +def test_atexit_rejects_direct_default_stack_creation(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + tracing_setup._shutdown_global_trace_provider() + + with pytest.raises(RuntimeError, match="Tracing shutdown has started"): + tracing_processors.default_exporter() + with pytest.raises(RuntimeError, match="Tracing shutdown has started"): + tracing_processors.default_processor() + + mock_client.assert_not_called() + + +@patch("httpx2.Client") +def test_atexit_closes_exporter_configured_before_bootstrap(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + + set_tracing_export_api_key("trace-only-key") + exporter = tracing_processors._global_exporter + tracing_setup._shutdown_global_trace_provider() + tracing_setup._shutdown_global_trace_provider() + + assert exporter is not None + assert exporter.is_shut_down is True + assert tracing_processors._global_processor is None + assert tracing_setup.GLOBAL_TRACE_PROVIDER is None + mock_client.return_value.close.assert_called_once() + + +@patch("httpx2.Client") +def test_atexit_closes_unattached_exporter_with_custom_provider(mock_client, monkeypatch): + _reset_default_tracing_stack(monkeypatch) + set_tracing_export_api_key("trace-only-key") + custom_provider = MagicMock() + tracing_setup.set_trace_provider(custom_provider) + + tracing_setup._shutdown_global_trace_provider() + + custom_provider.shutdown.assert_called_once() + mock_client.return_value.close.assert_called_once() + + @pytest.mark.serial @pytest.mark.review_optional def test_tracing_atexit_cleanup_timeout_preserves_process_exit_code_on_504() -> None: diff --git a/tests/tracing/test_setup.py b/tests/tracing/test_setup.py index a181f699bf..89ecc63269 100644 --- a/tests/tracing/test_setup.py +++ b/tests/tracing/test_setup.py @@ -41,6 +41,14 @@ def shutdown(self) -> None: self.shutdown_calls += 1 +@pytest.fixture(autouse=True) +def reset_atexit_shutdown_state(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(tracing_setup, "_ATEXIT_SHUTDOWN_STARTED", False) + monkeypatch.setattr(tracing_setup, "_SDK_DEFAULT_PROVIDER", None) + monkeypatch.setattr(tracing_setup, "_SDK_DEFAULT_PROCESSOR", None) + monkeypatch.setattr(tracing_processors, "_DEFAULT_STACK_ATEXIT_SHUTDOWN", False) + + def test_shutdown_global_trace_provider_calls_shutdown(monkeypatch: pytest.MonkeyPatch) -> None: provider = _DummyProvider() monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", provider)