fix(tracing): clear the exporter's shutdown request when a processor attaches

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._shutdown_lock = threading.Lock()
+        self._shutdown_requests = 0
 
         # Keep a client open for connection pooling across multiple export calls
         self._client = httpx2.Client(timeout=httpx2.Timeout(timeout=60, connect=5.0))
@@ -535,7 +537,24 @@ def close(self):
         self._client.close()
 
     def _request_shutdown(self) -> None:
-        self._shutdown_event.set()
+        with self._shutdown_lock:
+            self._shutdown_requests += 1
+            self._shutdown_event.set()
+
+    def _reset_shutdown(self) -> None:
+        """Drop one shutdown request, restoring retries once the last one is released.
+
+        Processors share an exporter -- `default_exporter()` hands the same instance to
+        every one of them -- so more than one can have a shutdown in flight. Retries come
+        back only when every request has been released; clearing on the first release would
+        resurrect the retry backoff of a worker still exporting under a shutdown of its own.
+        """
+        with self._shutdown_lock:
+            if self._shutdown_requests == 0:
+                return
+            self._shutdown_requests -= 1
+            if self._shutdown_requests == 0:
+                self._shutdown_event.clear()
 
 
 class BatchTraceProcessor(TracingProcessor):
@@ -580,6 +599,8 @@ def __init__(
         self._thread_start_lock = threading.Lock()
         self._export_lock = threading.Lock()
         self._shutdown_deadline: float | None = None
+        self._requested_exporter_shutdown = False
+        self._exporter_shutdown_lock = threading.Lock()
 
     def _ensure_thread_started(self) -> None:
         # Fast path without holding the lock
@@ -624,11 +645,10 @@ 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()
         if timeout is not None:
-            request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None)
-            if callable(request_exporter_shutdown):
-                request_exporter_shutdown()
+            self._request_exporter_shutdown()
+
+        self._shutdown_event.set()
 
         deadline = None if timeout is None else time.monotonic() + timeout
         self._shutdown_deadline = deadline
@@ -637,12 +657,57 @@ def shutdown(self, timeout: float | None = None):
         if self._worker_thread and self._worker_thread.is_alive():
             self._worker_thread.join(timeout=timeout)
             if self._worker_thread.is_alive():
+                # The worker outlived this shutdown, so it keeps the request until it stops.
                 logger.warning(
                     "[non-fatal] Tracing: shutdown timeout reached; dropping queued traces."
                 )
+            else:
+                # The worker released the request as it exited, unless it had already
+                # exited when we made it -- in which case this is the only release.
+                self._release_exporter_shutdown()
         else:
             # No background thread: process any remaining items synchronously.
             self._export_batches(deadline=deadline)
+            self._release_exporter_shutdown()
+
+    def _request_exporter_shutdown(self) -> None:
+        """Ask the exporter to abandon its retry backoff, at most once for this processor.
+
+        Under the lock so that concurrent `shutdown` calls make a single request between
+        them: the exporter counts requests, and this processor has exactly one release to
+        balance it with.
+        """
+        with self._exporter_shutdown_lock:
+            if self._requested_exporter_shutdown:
+                return
+            request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None)
+            if not callable(request_exporter_shutdown):
+                return
+            # Flagged before the request so a worker exiting concurrently either waits here
+            # and then releases what we asked for, or finds nothing to release yet and
+            # leaves it to the caller below, which releases once the worker has stopped.
+            self._requested_exporter_shutdown = True
+            request_exporter_shutdown()
+
+    def _release_exporter_shutdown(self) -> None:
+        """Give a shutdown request we made back to the exporter, now that our worker is done.
+
+        `default_exporter()` hands the same exporter to every processor, so leaving the
+        request outstanding makes the next processor give up on the first transient failure
+        instead of retrying -- blaming a shutdown that is long over. It is released only by
+        the processor that made it, and only once that processor's worker has stopped: while
+        the worker is still exporting it needs the cancellation to keep abandoning its
+        retries, including when `shutdown` timed out and returned without it. Other
+        processors' requests are counted separately by the exporter, so releasing ours never
+        takes the cancellation away from theirs.
+        """
+        with self._exporter_shutdown_lock:
+            if not self._requested_exporter_shutdown:
+                return
+            self._requested_exporter_shutdown = False
+            reset_exporter_shutdown = getattr(self._exporter, "_reset_shutdown", None)
+            if callable(reset_exporter_shutdown):
+                reset_exporter_shutdown()
 
     def force_flush(self):
         """
@@ -667,6 +732,9 @@ def _run(self):
         # Final drain after shutdown
         self._export_batches(deadline=self._shutdown_deadline)
 
+        # This worker is done exporting, so it no longer needs the exporter cancelled.
+        self._release_exporter_shutdown()
+
     def _export_batches(self, deadline: float | None = None):
         """Drains the queue and exports in batches of up to `max_batch_size` until the queue
         is completely empty.

tests/test_trace_processor.py

@@ -1,3 +1,4 @@
+import contextlib
 import logging
 import os
 import subprocess
@@ -612,6 +613,216 @@ def test_batch_trace_processor_shutdown_without_timeout_preserves_export_retries
     exporter.close()
 
 
+@patch("httpx2.Client")
+def test_new_processor_restores_exporter_retries_after_a_previous_shutdown(mock_client):
+    """A finished shutdown must not disable retries for the exporter's next processor.
+
+    ``_request_shutdown`` is one-way and ``default_exporter()`` hands the same instance to
+    every processor, so a stale request turned the first transient failure of the next
+    processor into an immediate give-up -- warning about a shutdown that was long over.
+    """
+    mock_response = MagicMock()
+    mock_response.status_code = 504
+    mock_client.return_value.post.return_value = mock_response
+
+    exporter = BackendSpanExporter(
+        api_key="test_key",
+        max_retries=3,
+        base_delay=0.01,
+        max_delay=0.02,
+    )
+
+    first = BatchTraceProcessor(exporter=exporter)
+    first.shutdown(timeout=1.0)
+    assert not exporter._shutdown_event.is_set()
+
+    second = BatchTraceProcessor(exporter=exporter)
+    second._queue.put_nowait(get_span(second))
+    second.force_flush()
+
+    assert mock_client.return_value.post.call_count == 3
+
+    exporter.close()
+
+
+@patch("httpx2.Client")
+def test_worker_that_outlives_shutdown_keeps_the_exporter_cancelled(mock_client):
+    """A worker abandoned by a timed-out shutdown owns the cancellation until it exits.
+
+    Releasing the request when the next processor attaches would hand it back while that
+    worker is still exporting, so it would sleep through its backoff and keep retrying after
+    the shutdown that was supposed to stop it had already returned.
+    """
+    post_called = threading.Event()
+    release_post = threading.Event()
+    mock_response = MagicMock()
+    mock_response.status_code = 504
+
+    def post(**kwargs: Any) -> Any:
+        post_called.set()
+        release_post.wait(timeout=5.0)
+        return mock_response
+
+    mock_client.return_value.post.side_effect = post
+
+    exporter = BackendSpanExporter(
+        api_key="test_key",
+        max_retries=100,
+        base_delay=10.0,
+        max_delay=10.0,
+    )
+    processor = BatchTraceProcessor(
+        exporter=exporter,
+        max_queue_size=1,
+        max_batch_size=1,
+        schedule_delay=60.0,
+        export_trigger_ratio=1.0,
+    )
+
+    processor.on_span_end(get_span(processor))
+    assert post_called.wait(timeout=2.0)
+
+    # The worker is stuck inside the request, so the join times out and shutdown returns
+    # with the worker still running and the exporter still cancelled.
+    processor.shutdown(timeout=0.05)
+    worker = processor._worker_thread
+    assert worker is not None and worker.is_alive()
+    assert exporter._shutdown_event.is_set()
+
+    # A replacement processor attaching must not take the cancellation away from it.
+    replacement = BatchTraceProcessor(exporter=exporter)
+    assert exporter._shutdown_event.is_set()
+
+    release_post.set()
+
+    # Still cancelled, so the abandoned worker abandons its backoff instead of sleeping
+    # through base_delay and retrying.
+    worker.join(timeout=2.0)
+    assert not worker.is_alive()
+    assert mock_client.return_value.post.call_count == 1
+
+    # And once it is gone it hands the exporter back, so the replacement can retry again.
+    assert not exporter._shutdown_event.is_set()
+
+    replacement.shutdown()
+    exporter.close()
+
+
+@patch("httpx2.Client")
+def test_shared_exporter_stays_cancelled_until_every_shutdown_releases(mock_client):
+    """Two processors can share one exporter, so one release must not speak for both.
+
+    A provider shuts its processors down in turn: the first times out with its worker still
+    inside a request, the second exits cleanly. If that second release cleared the shared
+    cancellation, the first worker would find it gone when its request finally returns a 5xx
+    and retry, after its own shutdown had already returned.
+    """
+    post_called = threading.Event()
+    release_post = threading.Event()
+    mock_response = MagicMock()
+    mock_response.status_code = 504
+
+    def post(**kwargs: Any) -> Any:
+        post_called.set()
+        release_post.wait(timeout=5.0)
+        return mock_response
+
+    mock_client.return_value.post.side_effect = post
+
+    exporter = BackendSpanExporter(
+        api_key="test_key",
+        max_retries=100,
+        base_delay=10.0,
+        max_delay=10.0,
+    )
+    blocked = BatchTraceProcessor(
+        exporter=exporter,
+        max_queue_size=1,
+        max_batch_size=1,
+        schedule_delay=60.0,
+        export_trigger_ratio=1.0,
+    )
+    other = BatchTraceProcessor(exporter=exporter)
+
+    blocked.on_span_end(get_span(blocked))
+    assert post_called.wait(timeout=2.0)
+
+    blocked.shutdown(timeout=0.05)
+    blocked_worker = blocked._worker_thread
+    assert blocked_worker is not None and blocked_worker.is_alive()
+
+    # The second processor never started a worker, so its shutdown finishes and releases
+    # right away -- while the first processor's request is still outstanding.
+    other.shutdown(timeout=1.0)
+    assert exporter._shutdown_event.is_set()
+
+    release_post.set()
+
+    blocked_worker.join(timeout=2.0)
+    assert not blocked_worker.is_alive()
+    assert mock_client.return_value.post.call_count == 1
+
+    # Both requests released now, so the exporter is usable again.
+    assert not exporter._shutdown_event.is_set()
+
+    exporter.close()
+
+
+@patch("httpx2.Client")
+def test_concurrent_shutdowns_leave_the_exporter_request_balanced(mock_client):
+    """`shutdown` is safe to call from several threads, so it must request only once.
+
+    The exporter counts outstanding requests and the worker releases one as it exits, so two
+    callers racing through the request path would leave a request outstanding forever -- and
+    every later processor giving up on its first transient failure instead of retrying.
+    """
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_client.return_value.post.return_value = mock_response
+
+    class RendezvousExporter(BackendSpanExporter):
+        """Blocks inside the window between the request check and the request itself.
+
+        `shutdown` reaches the exporter by looking `_request_shutdown` up on it, which
+        happens after a processor has decided to make the request and before it has recorded
+        that it did. Holding two callers here is what an unsynchronized decision would let
+        happen; serialized, only one of them ever arrives and the rendezvous times out.
+        """
+
+        def __init__(self, arrived: threading.Barrier, **kwargs: Any) -> None:
+            super().__init__(**kwargs)
+            self._arrived = arrived
+
+        @property
+        def _request_shutdown(self):
+            with contextlib.suppress(threading.BrokenBarrierError):
+                self._arrived.wait(timeout=0.5)
+            return super()._request_shutdown
+
+    exporter = RendezvousExporter(threading.Barrier(2), api_key="test_key")
+    processor = BatchTraceProcessor(exporter=exporter, schedule_delay=60.0)
+
+    # Start the worker, so the single release on its way out has to balance every request.
+    processor.on_span_end(get_span(processor))
+    assert processor._worker_thread is not None
+    assert processor._worker_thread.is_alive()
+
+    threads = [
+        threading.Thread(target=processor.shutdown, kwargs={"timeout": 2.0}) for _ in range(2)
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=5.0)
+        assert not thread.is_alive()
+
+    assert not processor._worker_thread.is_alive()
+    assert exporter._shutdown_requests == 0
+    assert not exporter._shutdown_event.is_set()
+
+    exporter.close()
+
+
 @pytest.mark.serial
 @pytest.mark.review_optional
 def test_tracing_atexit_cleanup_timeout_preserves_process_exit_code_on_504() -> None: