61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
|
"""In-process WebSocket connection registry for live server→client push.
|
||
|
|
|
||
|
|
Generic groundwork (notification module phase 2 is the first user): connections are kept in
|
||
|
|
this process's memory, keyed by user_id, so a request handler can push an event to all of a
|
||
|
|
user's open tabs/devices. Single-uvicorn-process assumption — if we ever run multiple workers,
|
||
|
|
replace the in-memory map with a Redis pub/sub fan-out behind the same `push()` interface (same
|
||
|
|
note as app.ratelimit).
|
||
|
|
|
||
|
|
`push()` is intentionally SYNC-callable: ordinary (sync) route handlers run in a threadpool, so
|
||
|
|
they schedule the actual send onto the WebSocket event loop via run_coroutine_threadsafe. The
|
||
|
|
connect/disconnect/_push coroutines all run on that one loop, so the connection map is only ever
|
||
|
|
mutated from a single thread.
|
||
|
|
"""
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
from starlette.websockets import WebSocket
|
||
|
|
|
||
|
|
log = logging.getLogger("siftlode.realtime")
|
||
|
|
|
||
|
|
|
||
|
|
class ConnectionManager:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self._conns: dict[int, set[WebSocket]] = defaultdict(set)
|
||
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
||
|
|
|
||
|
|
async def connect(self, user_id: int, ws: WebSocket) -> None:
|
||
|
|
self._loop = asyncio.get_running_loop()
|
||
|
|
await ws.accept()
|
||
|
|
self._conns[user_id].add(ws)
|
||
|
|
|
||
|
|
async def disconnect(self, user_id: int, ws: WebSocket) -> None:
|
||
|
|
conns = self._conns.get(user_id)
|
||
|
|
if conns:
|
||
|
|
conns.discard(ws)
|
||
|
|
if not conns:
|
||
|
|
self._conns.pop(user_id, None)
|
||
|
|
|
||
|
|
async def _push(self, user_id: int, payload: dict) -> None:
|
||
|
|
for ws in list(self._conns.get(user_id, ())):
|
||
|
|
try:
|
||
|
|
await ws.send_json(payload)
|
||
|
|
except Exception:
|
||
|
|
# Stale socket; it'll be cleaned up by its own disconnect handler.
|
||
|
|
pass
|
||
|
|
|
||
|
|
def push(self, user_id: int, payload: dict) -> None:
|
||
|
|
"""Schedule a push onto the WebSocket event loop. Safe to call from a sync route
|
||
|
|
handler (threadpool). No-op if no WebSocket has ever connected (no loop captured)."""
|
||
|
|
loop = self._loop
|
||
|
|
if loop is None:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
asyncio.run_coroutine_threadsafe(self._push(user_id, payload), loop)
|
||
|
|
except Exception:
|
||
|
|
log.debug("ws push failed", exc_info=True)
|
||
|
|
|
||
|
|
|
||
|
|
manager = ConnectionManager()
|