Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions tests/test_tui_gateway_ws.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import threading
import time

Expand Down Expand Up @@ -126,3 +127,40 @@ async def send_text(self, line):
loop.call_soon_threadsafe(loop.stop)
thread.join(timeout=2)
loop.close()


def test_ws_transport_serializes_concurrent_sends():
active_sends = 0
max_active_sends = 0
sent = []

class FakeWS:
async def send_text(self, line):
nonlocal active_sends, max_active_sends
active_sends += 1
max_active_sends = max(max_active_sends, active_sends)
try:
await asyncio.sleep(0.05)
sent.append(line)
finally:
active_sends -= 1

loop = asyncio.new_event_loop()
thread = threading.Thread(target=loop.run_forever, daemon=True)
thread.start()
try:
transport = ws_mod.WSTransport(FakeWS(), loop, peer="serialize-test")
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
futures = [
pool.submit(transport.write, {"idx": 1}),
pool.submit(transport.write, {"idx": 2}),
]
assert [f.result(timeout=2) for f in futures] == [True, True]

assert len(sent) == 2
assert max_active_sends == 1
assert transport._closed is False
finally:
loop.call_soon_threadsafe(loop.stop)
thread.join(timeout=2)
loop.close()
10 changes: 9 additions & 1 deletion tui_gateway/ws.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,11 @@ def __init__(
self._loop = loop
self._peer = peer
self._closed = False
# Starlette WebSocket sends are not a safe free-for-all: long foreground
# tasks can emit frames from several worker paths while the uvicorn loop
# is catching up. Serialize writes per socket so reconnect/log spam under
# high-output tasks is not amplified by overlapping send_text() calls.
self._send_lock = asyncio.Lock()

def write(self, obj: dict) -> bool:
if self._closed:
Expand Down Expand Up @@ -130,7 +135,10 @@ async def write_async(self, obj: dict) -> bool:

async def _safe_send(self, line: str) -> None:
try:
await self._ws.send_text(line)
async with self._send_lock:
if self._closed:
return
await self._ws.send_text(line)
except Exception as exc:
self._closed = True
_log.warning(
Expand Down