fix(ws): give the send deadline its own 30s clock and close the stalled socket
Follow-up on the cherry-picked #106372 fix (@gaoanze888), which reused the existing 10s _WS_WRITE_TIMEOUT_S as the deadline and only latched _closed. Why a separate, longer constant: _WS_WRITE_TIMEOUT_S is the worker-side wait for the LOOP to run the scheduled send; it deliberately does not mark the transport dead because a GIL-heavy turn stalls the loop for >10s routinely (#48445/#55545). The new clock is different in kind — asyncio.wait_for starts it only once the coroutine actually runs on the loop, and a healthy socket returns from send_text without waiting, so only real kernel backpressure can consume it (verified: a loop stall 3x longer than the deadline does not trip wait_for on 3.11/3.12/3.13 because the wakeup and the timer land in the same iteration and the result wins). Still, tying it to the 10s worker constant would let one test-time monkeypatch flip both semantics, and the issue asks explicitly not to declare the peer dead at the 10s worker wait. 30s is 3x the worker wait and under the client's 45s heartbeat deadline. Why close the socket: latching _closed alone leaves handle_ws parked in receive_text on a half-open peer, so the disconnect teardown (session detach/reap, live-transport unregister, client reconnect) never runs — the frontend keeps waiting on a socket that will never carry another frame. ws.close(1011) unblocks the read loop; websockets bounds the close itself (close_timeout -> abort) when the same full buffer stalls the close frame. Diagnostic now says which side stalled: "socket stalled, loop responsive" versus write()'s "loop stalled". Tests trimmed to the two invariants (≤2 per fix): a stalled send terminates within the deadline, fails both queued writes, latches closed and closes the socket; progress → complete → RPC reply ordering is preserved on a healthy socket. The stalled test uses the issue's own shape (two queued write_async calls) and monkeypatches the deadline with raising=False so on a base without the fix it fails on the symptom (sends never terminate), not on a missing attribute. Refs #106369
This commit is contained in:
teknium1
@@ -1,38 +1,34 @@
|
||||
"""A stalled ``send_text`` must not hold the writer lock forever (#106369).
|
||||
"""A stalled ``send_text`` must have a bounded lifetime (#106369).
|
||||
|
||||
Without a send deadline, socket backpressure parks ``_safe_send_many`` inside
|
||||
``_send_lock`` and ``_closed`` never latches, so reconnect recovery cannot
|
||||
start. With the fix, a stalled send times out, latches the transport closed,
|
||||
and a queued second batch returns immediately instead of blocking forever.
|
||||
``_safe_send_many`` awaits the socket while holding the connection-wide ``_send_lock``. Without a
|
||||
deadline, one send parked by socket backpressure trapped every later event and RPC reply behind the
|
||||
lock on an apparently open connection, so reconnect recovery never started.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import tui_gateway.ws as ws_mod
|
||||
from tui_gateway.ws import WSTransport
|
||||
|
||||
|
||||
class _StalledWS:
|
||||
"""``send_text`` never completes — models socket backpressure."""
|
||||
"""``send_text`` never completes (kernel backpressure); ``close`` is observable."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sent: list[str] = []
|
||||
self.closed_with: list[int] = []
|
||||
self._release = asyncio.Event()
|
||||
|
||||
async def send_text(self, line: str) -> None:
|
||||
await self._release.wait() # never set during the test
|
||||
self.sent.append(line)
|
||||
await self._release.wait()
|
||||
|
||||
async def close(self, code: int = 1000) -> None:
|
||||
self.closed_with.append(code)
|
||||
self._release.set()
|
||||
|
||||
|
||||
class _FastWS:
|
||||
"""``send_text`` completes immediately — guards against the timeout
|
||||
breaking the happy path."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sent: list[str] = []
|
||||
|
||||
@@ -40,54 +36,36 @@ class _FastWS:
|
||||
self.sent.append(line)
|
||||
|
||||
|
||||
def test_stalled_send_latches_closed() -> None:
|
||||
def test_stalled_send_closes_socket_and_releases_queued_reply(monkeypatch):
|
||||
# raising=False: on a base without the deadline the test must fail on the SYMPTOM (sends never terminate).
|
||||
monkeypatch.setattr("tui_gateway.ws._WS_SEND_DEADLINE_S", 0.05, raising=False)
|
||||
|
||||
async def _run() -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
ws = _StalledWS()
|
||||
transport = WSTransport(ws, loop, peer="127.0.0.1:1")
|
||||
orig = ws_mod._WS_WRITE_TIMEOUT_S
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = 0.05
|
||||
try:
|
||||
await transport._safe_send_many(["first"])
|
||||
finally:
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = orig
|
||||
assert transport.closed is True, "stalled send must latch the transport closed"
|
||||
assert ws.sent == [], "the stalled frame must not be recorded as sent"
|
||||
transport = WSTransport(ws, asyncio.get_running_loop(), peer="127.0.0.1:1")
|
||||
progress = asyncio.create_task(transport.write_async({"method": "event", "params": {"type": "tool.progress"}}))
|
||||
await asyncio.sleep(0) # progress is now inside send_text, holding _send_lock
|
||||
reply = asyncio.create_task(transport.write_async({"id": "submit", "result": {"status": "streaming"}}))
|
||||
# The loop stays responsive; both sends must still terminate within the deadline (not the 2s cap).
|
||||
results = await asyncio.wait_for(asyncio.gather(progress, reply), timeout=2.0)
|
||||
assert results == [False, False], "a send that missed the deadline must report failure, not success"
|
||||
assert transport.closed, "the transport must latch closed so handle_ws teardown/reconnect can run"
|
||||
await asyncio.sleep(0) # let the scheduled close task run
|
||||
assert ws.closed_with == [1011], "the stalled socket must be closed, not left half-open"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_queued_batch_returns_immediately_after_timeout() -> None:
|
||||
def test_progress_then_final_ordering_preserved_on_healthy_socket():
|
||||
async def _run() -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
ws = _StalledWS()
|
||||
transport = WSTransport(ws, loop, peer="127.0.0.1:1")
|
||||
orig = ws_mod._WS_WRITE_TIMEOUT_S
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = 0.05
|
||||
try:
|
||||
await transport._safe_send_many(["first"])
|
||||
# _closed is now True; the second batch must not queue on _send_lock.
|
||||
await transport._safe_send_many(["second"])
|
||||
finally:
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = orig
|
||||
assert transport.closed is True
|
||||
assert ws.sent == []
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_normal_send_completes_within_timeout() -> None:
|
||||
async def _run() -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
ws = _FastWS()
|
||||
transport = WSTransport(ws, loop, peer="127.0.0.1:1")
|
||||
orig = ws_mod._WS_WRITE_TIMEOUT_S
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = 0.05
|
||||
try:
|
||||
await transport._safe_send_many(["a", "b", "c"])
|
||||
finally:
|
||||
ws_mod._WS_WRITE_TIMEOUT_S = orig
|
||||
assert transport.closed is False, "a fast send must not latch the transport closed"
|
||||
assert ws.sent == ["a", "b", "c"]
|
||||
transport = WSTransport(ws, asyncio.get_running_loop(), peer="127.0.0.1:1")
|
||||
assert await transport.write_async({"method": "event", "params": {"type": "tool.progress"}})
|
||||
assert await transport.write_async({"method": "event", "params": {"type": "message.complete"}})
|
||||
assert await transport.write_async({"id": "submit", "result": {"status": "done"}})
|
||||
assert not transport.closed
|
||||
assert [json.loads(s).get("params", {}).get("type", "reply") for s in ws.sent] == [
|
||||
"tool.progress", "message.complete", "reply",
|
||||
]
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
@@ -58,6 +58,13 @@ def _sanitize_ws_text(text: str) -> str:
|
||||
# Max seconds a pool-dispatched handler blocks waiting for the loop to flush a WS frame before we
|
||||
# give up waiting (the transport is NOT marked dead).
|
||||
_WS_WRITE_TIMEOUT_S = 10.0
|
||||
# Max seconds one send_text may await the socket once it is actually running on the loop. A healthy
|
||||
# socket returns from send_text without waiting (the frame lands in the transport buffer); only kernel
|
||||
# backpressure parks it, so a GIL/loop stall cannot start this clock. Deliberately 3x the worker wait
|
||||
# above and under the client's 45s heartbeat deadline (apps/shared json-rpc-gateway): a peer that
|
||||
# cannot drain ~48 KiB in 30s is gone, and closing here starts its reconnect instead of leaving every
|
||||
# later frame and RPC reply parked behind the writer lock (#106369).
|
||||
_WS_SEND_DEADLINE_S = 30.0
|
||||
_WS_LOG_PAYLOAD_PREVIEW = 240
|
||||
|
||||
# Per-token streaming frames are coalesced: buffered and flushed as a batch on a short timer instead
|
||||
@@ -183,14 +190,16 @@ class WSTransport:
|
||||
return
|
||||
payload = _sanitize_ws_text(line)
|
||||
try:
|
||||
await asyncio.wait_for(self._ws.send_text(payload), timeout=_WS_WRITE_TIMEOUT_S)
|
||||
except asyncio.TimeoutError as exc:
|
||||
# A stalled send_text (socket backpressure) must not hold the writer lock forever.
|
||||
# Unlike the loop-stall case in write(), this means the socket itself is unresponsive:
|
||||
# latch closed so queued batches bail and reconnect recovery can start. See #106369.
|
||||
await asyncio.wait_for(self._ws.send_text(payload), timeout=_WS_SEND_DEADLINE_S)
|
||||
except asyncio.TimeoutError:
|
||||
# The loop is responsive (the timer fired) but the socket never drained: unlike the
|
||||
# loop-stall wait in write(), this is a dead peer. Latch under the writer lock so queued
|
||||
# batches bail, and close the socket so handle_ws's read loop ends and its teardown
|
||||
# (session detach/reap, client reconnect) runs. See #106369.
|
||||
self._closed = True
|
||||
_log.warning("ws send timed out peer=%s timeout=%ss error_type=%s error=%s",
|
||||
self._peer, _WS_WRITE_TIMEOUT_S, type(exc).__name__, exc)
|
||||
_log.warning("ws send deadline exceeded (socket stalled, loop responsive) peer=%s deadline=%ss — closing",
|
||||
self._peer, _WS_SEND_DEADLINE_S)
|
||||
self._loop.create_task(self._close_stalled_socket())
|
||||
return
|
||||
except UnicodeEncodeError as exc:
|
||||
# A single illegal UTF-8 frame (lone surrogate) must not tear down the socket.
|
||||
@@ -208,6 +217,14 @@ class WSTransport:
|
||||
self._token_flush_handle.cancel()
|
||||
self._token_flush_handle = None
|
||||
|
||||
async def _close_stalled_socket(self) -> None:
|
||||
"""Close the peer socket after a send deadline so ``handle_ws``'s ``receive_text`` unblocks and its
|
||||
disconnect teardown runs. The server library bounds this (websockets ``close_timeout`` → abort)."""
|
||||
try:
|
||||
await self._ws.close(code=1011)
|
||||
except Exception as exc: # noqa: BLE001 - the peer is already gone; teardown is what matters
|
||||
_log.debug("ws close after send deadline failed peer=%s error=%s", self._peer, exc)
|
||||
|
||||
|
||||
def _ws_peer_label(ws: Any) -> str:
|
||||
"""``host:port`` when available, else a stable placeholder."""
|
||||
|
||||
Reference in New Issue
Block a user