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-scheduled close task now consults the evicted-client closer first:
a busy client (connection in flight) is queued and closed by the closer once idle
and out of grace (EvictedClientCloser.defer_if_busy), an idle one is closed in
the task as before — the existing task-reference machinery, no-loop and
cross-loop disposal branches are unchanged. The in-flight check runs inside the
task, not in __del__, so no closer 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-09-10 18:05:26 +00:00
parent 9a715df212
commit 389c4befca
4 changed files with 153 additions and 1 deletions

View file

@ -206,6 +206,36 @@ class EvictedClientCloser:
)
)
def defer_if_busy(self, client: object) -> bool:
"""Queue a deferred close for a client with a request in flight; no-op when idle.
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. A busy client is
queued and closed by ``reap`` once idle and out of grace; returns True when the
close was deferred (the caller must NOT close the client itself), False when the
client is idle and the caller may close it directly.
"""
if _close_function(client) is None:
return False
if not _has_connection_in_flight(client):
return False
self.mark_owned(client)
self.schedule(client)
return True
def close_or_defer(self, client: object) -> None:
"""Close an unreferenced client now if idle, else queue it for a deferred close.
Same in-flight rule as ``defer_if_busy``; an idle client is closed immediately,
preserving the reclamation the finalizer used to do.
"""
if _close_function(client) is None:
return
if self.defer_if_busy(client):
return
self._close(client)
def reap(self) -> None:
"""Close every queued client that is due, idle, and closable from here.

View file

@ -1102,13 +1102,34 @@ 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())
task: Final = loop.create_task(self._finalizer_close_client(self._client))
cls: Final = type(self)
cls._finalizer_close_tasks.add(task)
task.add_done_callback(cls._on_finalizer_close_done)
except Exception:
pass
@staticmethod
async def _finalizer_close_client(client: httpx.AsyncClient) -> None:
"""Close a finalized handler's client, unless a request is still in flight.
The sole-referrer refcount guard in ``__del__`` proves nothing else holds the
client *object*, but a request in flight references only the pooled connection,
so it is invisible to that check: closing here would tear the pool down under
every live SSE stream (a cache-evicted handler is finalized the moment the cache
drops it one batch of mid-turn stream deaths per handler-cache TTL per
process). A busy client is handed to the evicted-client closer, which closes it
once it reports no connection in flight and a grace window has passed; an idle
one is closed here so the finalizer keeps its reclamation. The in-flight check
runs inside this task, not in ``__del__``, so no closer lock is taken in GC
context.
"""
from litellm.caching.evicted_client_closer import default_evicted_client_closer
if default_evicted_client_closer.defer_if_busy(client):
return
await client.aclose()
@staticmethod
def _create_async_transport(
ssl_context: ssl.SSLContext | None = None,

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

@ -1675,3 +1675,50 @@ async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch
finally:
await handler.close()
assert closed.is_set()
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()