Optimize websocket broadcast fanout

hub/ws/manager.py

@@ -3,23 +3,26 @@
 import asyncio
 import json
 from collections import defaultdict
+from typing import Protocol
 
-from fastapi import WebSocket
 
-_subscribers: dict[str, set[WebSocket]] = defaultdict(set)
+class WebSocketLike(Protocol):
+    async def send_text(self, data: str) -> None: ...
+
+_subscribers: dict[str, set[WebSocketLike]] = defaultdict(set)
 _lock = asyncio.Lock()
 
 # Price feed subscribers (not agent-specific)
-_price_subscribers: set[WebSocket] = set()
+_price_subscribers: set[WebSocketLike] = set()
 _price_lock = asyncio.Lock()
 
 
-async def subscribe(agent_id: str, ws: WebSocket):
+async def subscribe(agent_id: str, ws: WebSocketLike):
     async with _lock:
         _subscribers[agent_id].add(ws)
 
 
-async def unsubscribe(agent_id: str, ws: WebSocket):
+async def unsubscribe(agent_id: str, ws: WebSocketLike):
     async with _lock:
         _subscribers[agent_id].discard(ws)
         if not _subscribers[agent_id]:
@@ -30,23 +33,24 @@ async def broadcast(agent_id: str, msg_type: str, data):
     """Send to all subscribers of an agent."""
     async with _lock:
         subs = list(_subscribers.get(agent_id, []))
+    if not subs:
+        return
     payload = json.dumps({"type": msg_type, "data": data})
-    dead = []
-    for ws in subs:
-        try:
-            await ws.send_text(payload)
-        except Exception:
-            dead.append(ws)
+    results = await asyncio.gather(
+        *(ws.send_text(payload) for ws in subs),
+        return_exceptions=True,
+    )
+    dead = [ws for ws, result in zip(subs, results) if isinstance(result, Exception)]
     for ws in dead:
         await unsubscribe(agent_id, ws)
 
 
-async def subscribe_prices(ws: WebSocket):
+async def subscribe_prices(ws: WebSocketLike):
     async with _price_lock:
         _price_subscribers.add(ws)
 
 
-async def unsubscribe_prices(ws: WebSocket):
+async def unsubscribe_prices(ws: WebSocketLike):
     async with _price_lock:
         _price_subscribers.discard(ws)
 
@@ -58,11 +62,10 @@ async def broadcast_prices(price_cache: dict):
     if not subs:
         return
     payload = json.dumps({"type": "prices", "data": price_cache})
-    dead = []
-    for ws in subs:
-        try:
-            await ws.send_text(payload)
-        except Exception:
-            dead.append(ws)
+    results = await asyncio.gather(
+        *(ws.send_text(payload) for ws in subs),
+        return_exceptions=True,
+    )
+    dead = [ws for ws, result in zip(subs, results) if isinstance(result, Exception)]
     for ws in dead:
         await unsubscribe_prices(ws)

tests/test_ws_manager.py

@@ -0,0 +1,60 @@
+import asyncio
+import time
+
+from hub.ws import manager
+
+
+class FakeWebSocket:
+    def __init__(self, *, delay: float = 0.0, fail: bool = False):
+        self.delay = delay
+        self.fail = fail
+        self.messages: list[str] = []
+
+    async def send_text(self, payload: str) -> None:
+        await asyncio.sleep(self.delay)
+        if self.fail:
+            raise RuntimeError("closed")
+        self.messages.append(payload)
+
+
+def clear_subscribers() -> None:
+    manager._subscribers.clear()
+    manager._price_subscribers.clear()
+
+
+def test_agent_broadcast_sends_to_slow_subscribers_concurrently():
+    async def run() -> None:
+        clear_subscribers()
+        sockets = [FakeWebSocket(delay=0.03) for _ in range(8)]
+        for ws in sockets:
+            await manager.subscribe("agent-1", ws)
+
+        start = time.perf_counter()
+        await manager.broadcast("agent-1", "status", {"ok": True})
+        elapsed = time.perf_counter() - start
+
+        assert elapsed < 0.12
+        assert all(len(ws.messages) == 1 for ws in sockets)
+        clear_subscribers()
+
+    asyncio.run(run())
+
+
+def test_price_broadcast_sends_to_slow_subscribers_concurrently_and_removes_dead():
+    async def run() -> None:
+        clear_subscribers()
+        live = [FakeWebSocket(delay=0.03) for _ in range(6)]
+        dead = FakeWebSocket(delay=0.03, fail=True)
+        for ws in [*live, dead]:
+            await manager.subscribe_prices(ws)
+
+        start = time.perf_counter()
+        await manager.broadcast_prices({"BTCUSDT": {"price": 100.0}})
+        elapsed = time.perf_counter() - start
+
+        assert elapsed < 0.12
+        assert all(len(ws.messages) == 1 for ws in live)
+        assert dead not in manager._price_subscribers
+        clear_subscribers()
+
+    asyncio.run(run())