fix(http_handler): defer finalizer client close while requests are in flight

A cache-evicted AsyncHTTPHandler is finalized the moment the cache drops it. The
finalizer's sole-referrer refcount guard proves nothing else holds the client
object, but a request in flight references only the pooled connection, so the
guard cannot see it: closing there tears the pool down under every live SSE
stream, one batch per handler-cache TTL per process. The finalizer now defers to
EvictedClientCloser.close_or_defer, which closes an idle client immediately and
queues a busy one until it reports no connection in flight and the grace window
has passed. The handoff runs via loop.call_soon so no lock is taken in GC
context.

Resolves the hourly batched mid-stream stream deaths observed in a production
proxy deployment (streams died in same-second batches at a fixed phase after
process start, matching the client-cache TTL).
This commit is contained in:
Amin Saedi 2026-08-25 13:44:25 +00:00 committed by Amin Saedi
parent 5e4b3838aa
commit 36104956c0
4 changed files with 171 additions and 58 deletions

View file

@ -206,6 +206,23 @@ class EvictedClientCloser:
)
)
def close_or_defer(self, client: object) -> None:
"""Close an unreferenced client now if idle, else queue it for a deferred close.
For a finalized handler's client the sole-referrer refcount check has already
proven nothing else holds the client object, but a request in flight references
only the pooled connection, so it is invisible to that check. An idle client is
closed immediately, preserving the reclamation the finalizer used to do; a busy
one is queued and closed by ``reap`` once idle and out of grace.
"""
if _close_function(client) is None:
return
if not _has_connection_in_flight(client):
self._close(client)
return
self.mark_owned(client)
self.schedule(client)
def reap(self) -> None:
"""Close every queued client that is due, idle, and closable from here.

View file

@ -9,7 +9,7 @@ import threading
import time
from collections.abc import AsyncIterable, Callable, Iterable, Mapping
from http.cookiejar import CookieJar, DefaultCookiePolicy
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict
from typing import TYPE_CHECKING, Any, Final, NoReturn, Optional, TypeAlias, TypedDict
import certifi
import httpx
@ -160,6 +160,26 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool:
return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER
def _hand_off_client_to_evicted_closer(client: object) -> None:
"""
Route a finalized handler's client to the evicted-client closer instead of closing it.
A request in flight holds references to the pooled connection, not to the client
object, so the sole-referrer refcount check cannot see it. Closing from ``__del__``
therefore tears the pool down under live requests: a cache-evicted handler is
finalized the moment the cache drops it, and every SSE stream on its pool dies
mid-turn (one batch per handler-cache TTL per process). The closer closes an idle
client immediately (preserving reclamation) and defers a busy one until it reports
no connection in flight and a grace window has passed.
Runs via ``loop.call_soon`` rather than inside the finalizer, so no lock is taken
in GC context.
"""
from litellm.caching.evicted_client_closer import default_evicted_client_closer
default_evicted_client_closer.close_or_defer(client)
def blocked_cookie_jar() -> CookieJar:
"""A jar that stores no response cookie and sends none, for httpx clients.
@ -933,21 +953,6 @@ class AsyncHTTPHandler:
response.raise_for_status()
return response
# Strong references to finalizer-scheduled client-close tasks. A bare
# create_task() result may be garbage-collected before it runs, leaving
# the underlying aiohttp session unclosed ("Unclosed client session").
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
@classmethod
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
cls._finalizer_close_tasks.discard(task)
if task.cancelled():
return
exc: Final = task.exception()
if exc is not None:
verbose_logger.debug("Error closing client at finalization: %s", exc)
def _aiohttp_session_bound_elsewhere(self, loop: asyncio.AbstractEventLoop) -> bool:
"""True when the wrapped aiohttp session is bound to a loop other than
``loop`` awaiting ``aclose()`` here would touch that loop's internals."""
@ -1006,10 +1011,7 @@ class AsyncHTTPHandler:
# here is the cross-loop path the transport refuses.
self._dispose_wrapped_aiohttp_session()
return
task: Final = loop.create_task(self._client.aclose())
cls: Final = type(self)
cls._finalizer_close_tasks.add(task)
task.add_done_callback(cls._on_finalizer_close_done)
loop.call_soon(_hand_off_client_to_evicted_closer, self._client)
except Exception:
pass

View file

@ -409,3 +409,57 @@ def test_a_reap_looks_at_what_is_due_rather_than_at_the_whole_queue():
f"{CountingDeadline.comparisons} deadline comparisons for {evictions} evictions; "
"a reap is walking the whole queue"
)
@pytest.mark.asyncio
async def test_close_or_defer_closes_an_idle_client_immediately():
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.close_or_defer(client)
await asyncio.sleep(0.05)
assert client.closed is True
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_close_or_defer_defers_a_client_with_a_request_on_the_wire():
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
client = httpx.AsyncClient()
async with asyncio.timeout(30):
async with client.stream("GET", f"http://127.0.0.1:{port}/") as response:
body_iter = response.aiter_raw()
await body_iter.__anext__()
closer.close_or_defer(client)
assert not client.is_closed
assert closer.pending_count == 1
remainder = b"".join([chunk async for chunk in body_iter])
assert b"hello" in remainder
clock.advance(61.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed
assert closer.pending_count == 0
# No wait_closed(): on Python >= 3.12.1 it waits for every client
# transport, and a pooled keepalive connection would park it forever.
server.close()
def test_close_or_defer_ignores_a_value_without_a_close_function():
clock = FakeClock()
closer = make_closer(clock)
closer.close_or_defer(object())
assert closer.pending_count == 0

View file

@ -1,8 +1,6 @@
import asyncio
import gc
import io
import os
import pathlib
import ssl
import threading
import weakref
@ -364,9 +362,11 @@ async def test_async_handler_with_shared_session():
async def test_get_async_httpx_client_with_shared_session():
"""Test get_async_httpx_client with shared session"""
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
AsyncHTTPHandler as AsyncHTTPHandlerReload,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.types.utils import LlmProviders
# Create a mock shared session
@ -388,9 +388,11 @@ async def test_get_async_httpx_client_with_shared_session():
async def test_get_async_httpx_client_without_shared_session():
"""Test get_async_httpx_client without shared session (backward compatibility)"""
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
AsyncHTTPHandler as AsyncHTTPHandlerReload,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.types.utils import LlmProviders
# Test without shared session
@ -426,6 +428,7 @@ async def test_session_reuse_chain():
def test_shared_session_parameter_in_acompletion():
"""Test that acompletion function accepts shared_session parameter"""
import inspect
from litellm.main import acompletion
# Get the function signature
@ -443,6 +446,7 @@ def test_shared_session_parameter_in_acompletion():
def test_shared_session_parameter_in_completion():
"""Test that completion function accepts shared_session parameter"""
import inspect
from litellm.main import completion
# Get the function signature
@ -461,9 +465,11 @@ def test_shared_session_parameter_in_completion():
async def test_session_reuse_integration():
"""Integration test for session reuse functionality"""
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
AsyncHTTPHandler as AsyncHTTPHandlerReload,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.types.utils import LlmProviders
# Create a mock session
@ -1225,10 +1231,13 @@ def test_finalizer_without_running_loop_closes_dead_loop_session():
@pytest.mark.asyncio
async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
"""With a running loop, finalization schedules an async close and must keep
a strong reference to the task until it completes a bare create_task()
result may be collected before it runs, leaving the session unclosed."""
async def test_finalizer_with_running_loop_hands_idle_client_to_closer():
"""With a running loop, finalization defers to the evicted-client closer,
which closes an idle client immediately while holding a strong reference to
the close task a bare create_task() result may be collected before it
runs, leaving the session unclosed."""
from litellm.caching.evicted_client_closer import default_evicted_client_closer
handler = AsyncHTTPHandler(timeout=61.0)
transport = handler.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
@ -1236,16 +1245,17 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
assert not session.closed
del transport
baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks)
baseline_tasks = set(default_evicted_client_closer._close_tasks)
del handler
gc.collect()
await asyncio.sleep(0)
scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks
scheduled = default_evicted_client_closer._close_tasks - baseline_tasks
assert len(scheduled) == 1
await asyncio.gather(*scheduled)
assert session.closed
assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
assert not (default_evicted_client_closer._close_tasks & scheduled)
@pytest.mark.asyncio
@ -1277,40 +1287,70 @@ async def test_sync_close_helper_respects_session_ownership():
await owned_handler.close()
@pytest.mark.asyncio
async def test_finalizer_close_done_consumes_exception():
"""A failing finalizer close must have its exception retrieved by the done
callback, or asyncio emits "Task exception was never retrieved" at GC
the same log noise the finalizer path exists to eliminate."""
async def failing_close() -> None:
raise RuntimeError("close failed")
task = asyncio.get_running_loop().create_task(failing_close())
AsyncHTTPHandler._finalizer_close_tasks.add(task)
await asyncio.sleep(0)
AsyncHTTPHandler._on_finalizer_close_done(task)
assert task not in AsyncHTTPHandler._finalizer_close_tasks
cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30))
cancelled.cancel()
await asyncio.sleep(0)
AsyncHTTPHandler._on_finalizer_close_done(cancelled)
@pytest.mark.asyncio
async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling():
"""GC on a live loop (e.g. the app's) of a handler whose session belongs to
another, dead loop must not schedule aclose() here that is the cross-loop
path the transport refuses and must still dispose the session."""
another, dead loop must not hand the client to the closer that is the
cross-loop path the transport refuses and must still dispose the session."""
from litellm.caching.evicted_client_closer import default_evicted_client_closer
handler = AsyncHTTPHandler(timeout=61.0)
session = await asyncio.to_thread(_mint_session_on_dead_loop, handler)
assert not session.closed
baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks)
baseline_tasks = set(default_evicted_client_closer._close_tasks)
baseline_pending = default_evicted_client_closer.pending_count
del handler
gc.collect()
await asyncio.sleep(0)
assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks
assert default_evicted_client_closer._close_tasks == baseline_tasks
assert default_evicted_client_closer.pending_count == baseline_pending
assert session.closed
async def _slow_chunked_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
"""Serves a chunked body in two installments, so a stream is on the wire while the handler dies."""
await reader.read(4096)
writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
writer.write(b"5\r\nfirst\r\n")
await writer.drain()
await asyncio.sleep(0.4)
writer.write(b"4\r\nlast\r\n0\r\n\r\n")
await writer.drain()
@pytest.mark.asyncio
async def test_collected_handler_never_kills_a_stream_in_flight(monkeypatch):
"""
Regression: a cache-evicted (hence collected) handler's finalizer used to close the
owned client while its pool still served live SSE streams, killing every one of them
mid-turn once per handler-cache TTL. The finalizer must defer to the evicted-client
closer while a connection is in flight, so the stream reads to completion.
"""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "force_ipv4", False)
server = await asyncio.start_server(_slow_chunked_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
handler = AsyncHTTPHandler()
async with asyncio.timeout(30):
request = handler.client.build_request("GET", f"http://127.0.0.1:{port}/")
response = await handler.client.send(request, stream=True)
body_iter = response.aiter_raw()
first = await body_iter.__anext__()
assert b"first" in first
del handler
gc.collect()
await asyncio.sleep(0.1)
remainder = b"".join([chunk async for chunk in body_iter])
assert b"last" in remainder
await response.aclose()
# No wait_closed(): the surviving client's pooled keepalive connection is
# the point of this test, and on Python >= 3.12.1 wait_closed() waits for
# every client transport, parking the suite forever.
server.close()