Reduce thread wakeups by using proper async

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:

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