diff --git a/modalapi/websocket_bridge.py b/modalapi/websocket_bridge.py index 9902f7bb1..a464dedda 100644 --- a/modalapi/websocket_bridge.py +++ b/modalapi/websocket_bridge.py @@ -55,6 +55,7 @@ def __init__( self.ws = None self._loop: Optional[asyncio.AbstractEventLoop] = None self._stop_event: asyncio.Event = asyncio.Event() + self._wakeup: asyncio.Event = asyncio.Event() # Metrics self.messages_sent = 0 @@ -80,10 +81,21 @@ def signal_stop(self): if self._loop is None: return self._loop.call_soon_threadsafe(self._stop_event.set) + self._loop.call_soon_threadsafe(self._wakeup.set) ws = self.ws if ws is not None: asyncio.run_coroutine_threadsafe(ws.close(), self._loop) + def notify(self): + """Thread-safe: wake the send loop after a message is enqueued.""" + loop = self._loop + if loop is None: + return # worker not started; connect-time flush covers these + try: + loop.call_soon_threadsafe(self._wakeup.set) + except RuntimeError: + pass # loop closed during shutdown + async def _interruptible_sleep(self, delay: float) -> bool: """Sleep for delay seconds; returns True if stop was signaled before the delay elapsed.""" try: @@ -124,7 +136,19 @@ async def _async_worker(self): if flushed: logging.info(f"Flushed {flushed} stale messages from queue after reconnect") - await asyncio.gather(self._process_queue(ws), self._receive_messages(ws)) + # FIRST_COMPLETED, not gather: the send loop parks on _wakeup and + # cannot notice a closed socket on its own. Whichever loop exits + # first cancels the other so we fall through to reconnect. + tasks = { + asyncio.create_task(self._process_queue(ws)), + asyncio.create_task(self._receive_messages(ws)), + } + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + for task in done: + task.result() # re-raise so the reconnect handler sees it except (websockets.exceptions.WebSocketException, OSError, ConnectionRefusedError) as e: logging.error(f"WebSocket connection error: {e}") @@ -156,7 +180,11 @@ async def _process_queue(self, ws): try: msg = self.command_queue.get_nowait() except queue.Empty: - await asyncio.sleep(0.001) # 1ms yield + # Clear before re-checking: a producer that enqueues between the + # failed get and the clear would otherwise have its wakeup erased. + self._wakeup.clear() + if self.command_queue.empty(): + await self._wakeup.wait() continue await ws.send(msg) @@ -263,6 +291,7 @@ def send_bpm(self, bpm: float) -> bool: if self._worker.backpressure_active: return False self.command_queue.put_nowait(f"transport-bpm {bpm}") + self._worker.notify() return True def send_parameter(self, instance_id: str, symbol: str, value: float) -> bool: @@ -271,6 +300,7 @@ def send_parameter(self, instance_id: str, symbol: str, value: float) -> bool: if self._worker.backpressure_active: return False self.command_queue.put_nowait(f"param_set /graph/{instance_id}/{symbol} {value}") + self._worker.notify() return True def get_received_messages(self) -> list: diff --git a/tests/test_websocket_bridge.py b/tests/test_websocket_bridge.py index ff10395fa..5ea828ab6 100644 --- a/tests/test_websocket_bridge.py +++ b/tests/test_websocket_bridge.py @@ -3,6 +3,8 @@ import asyncio import queue +import pytest + from modalapi.websocket_bridge import AsyncWebSocketBridge, WebSocketWorker @@ -231,3 +233,82 @@ def test_multiple_sends_preserve_order(): "transport-bpm 60", "param_set /graph/b/y 2.0", ] + + +# --------------------------------------------------------------------------- +# _process_queue: parks on _wakeup rather than polling +# --------------------------------------------------------------------------- + + +class _SendWs: + """WebSocket stand-in that records sends. No transport => buffer size reads as 0.""" + + def __init__(self): + self.sent: list[str] = [] + + async def send(self, msg: str) -> None: + self.sent.append(msg) + + +async def _started_worker() -> tuple[WebSocketWorker, _SendWs, asyncio.Task]: + worker = _make_worker() + worker.running = True + worker._loop = asyncio.get_running_loop() + ws = _SendWs() + task = asyncio.create_task(worker._process_queue(ws)) + await asyncio.sleep(0.02) # let it reach the park + return worker, ws, task + + +def test_idle_send_loop_parks_instead_of_polling(): + async def go(): + worker, ws, task = await _started_worker() + assert not task.done() + assert ws.sent == [] + # Parked: the event was consumed, so the loop is suspended, not spinning. + assert not worker._wakeup.is_set() + + worker.running = False + worker.signal_stop() + await asyncio.wait_for(task, timeout=1.0) + + asyncio.run(go()) + + +def test_notify_wakes_parked_send_loop(): + async def go(): + worker, ws, task = await _started_worker() + + worker.command_queue.put_nowait("param_set /graph/a/x 1.0") + worker.notify() + await asyncio.sleep(0.02) + assert ws.sent == ["param_set /graph/a/x 1.0"] + + worker.running = False + worker.signal_stop() + await asyncio.wait_for(task, timeout=1.0) + + asyncio.run(go()) + + +def test_parked_send_loop_is_cancellable(): + """_async_worker cancels the sender when the receive loop sees the socket close. + + Without this, a disconnect while idle would hang the bridge forever: the sender + never touches the socket, so it cannot notice the close on its own. + """ + + async def go(): + _worker, _ws, task = await _started_worker() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(go()) + + +def test_notify_before_worker_starts_is_a_noop(): + # No event loop yet; notify() must not raise. Queued messages are flushed on connect. + bridge = _make_bridge() + bridge.send_parameter("a", "x", 1.0) + assert bridge.get_queue_depth() == 1