diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index b60bcde2b5a..bb1abb493db 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -430,6 +430,7 @@ class TenantFanOutSpanProcessor(SpanProcessor): shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS, operator_sinks: frozenset[_SinkKey] = frozenset(), pending_drains: int = _MAX_PENDING_DRAINS, + drain_pool: _DrainPool | None = None, ) -> None: self._operator_sinks: Final = operator_sinks self._drain_seconds: Final = shutdown_drain_seconds @@ -439,7 +440,7 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count - self._drain: Final = _DrainPool(capacity=pending_drains) + self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains) def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: return None @@ -537,9 +538,9 @@ class TenantFanOutSpanProcessor(SpanProcessor): return False built: Final = self._cached_or_built_locked(destination) drained: Final = self._drainable_locked() - for shed in drained: - self._drain.submit(shed) - return built is not None + for shed in drained: + self._drain.submit(shed) + return built is not None def _acquire(self, destination: "OtelDestination") -> SpanProcessor | None: """The processor for ``destination``, marked busy until ``_release``. @@ -557,9 +558,9 @@ class TenantFanOutSpanProcessor(SpanProcessor): return None self._exporting[id(processor)] = self._exporting.get(id(processor), 0) + 1 drained: Final = self._drainable_locked() - for shed in drained: - self._drain.submit(shed) - return processor + for shed in drained: + self._drain.submit(shed) + return processor def _cached_or_built_locked(self, destination: "OtelDestination") -> SpanProcessor | None: key: Final = destination.cache_key() @@ -596,8 +597,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): if not self._exporting: self._lock.notify_all() drained: Final = self._drainable_locked() - for retired in drained: - self._drain.submit(retired) + for retired in drained: + self._drain.submit(retired) def _retire_overflow_locked(self) -> None: """Move the LRU processor out of the cache once it is past the cap.""" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py index 4cdc1d93e78..410c4ae4cc6 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py @@ -1735,6 +1735,58 @@ class TestEvictionSafety: time.sleep(0.02) assert fan_out.deliverable((self._dest(999),)) == (self._dest(999),), "the fan-out never recovered" + + def test_concurrent_eviction_cannot_build_between_retirement_and_drain_submission(self): + """A second request cannot build while the first eviction is being handed to + the drain, or concurrent churn can outrun the pending-drain limit.""" + import threading + + from litellm.integrations.otel.plumbing.providers import ( + _DrainPool, + _MAX_CACHED_DESTINATION_PROCESSORS, + ) + + class GatedDrain(_DrainPool): + def __init__(self): + super().__init__(workers=0) + self.started = threading.Event() + self.release = threading.Event() + + def saturated(self): + return False + + def submit(self, processor): + if not self.started.is_set(): + self.started.set() + self.release.wait(timeout=5) + + built = [] + + def factory(_destination): + built.append(self.Recording()) + return built[-1] + + drain = GatedDrain() + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, drain_pool=drain) + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS): + fan_out._release(fan_out._acquire(self._dest(index))) + + first = threading.Thread(target=lambda: fan_out._release(fan_out._acquire(self._dest(32)))) + first.start() + assert drain.started.wait(timeout=5) + second = threading.Thread(target=lambda: fan_out._release(fan_out._acquire(self._dest(33)))) + second.start() + time.sleep(0.1) + + assert len(built) == _MAX_CACHED_DESTINATION_PROCESSORS + 1 + + drain.release.set() + first.join(timeout=5) + second.join(timeout=5) + assert not first.is_alive() and not second.is_alive() + assert len(built) == _MAX_CACHED_DESTINATION_PROCESSORS + 2 + + def test_drain_workers_are_daemons(self): """Python joins a ThreadPoolExecutor's workers at interpreter exit, so one unreachable tenant collector would hold the proxy open for its export timeout on the way down."""