Initial commit: stock market platform

backend/app/websocket/manager.py

@@ -0,0 +1,58 @@
+"""
+WebSocket connection manager — broadcasts real-time stock quotes to subscribers.
+"""
+import asyncio
+import json
+from collections import defaultdict
+from fastapi import WebSocket
+from loguru import logger
+
+
+class ConnectionManager:
+    def __init__(self):
+        # symbol -> set of WebSocket connections
+        self._subs: dict[str, set[WebSocket]] = defaultdict(set)
+        self._all: set[WebSocket] = set()
+
+    async def connect(self, ws: WebSocket, symbol: str | None = None):
+        await ws.accept()
+        self._all.add(ws)
+        if symbol:
+            self._subs[symbol].add(ws)
+        logger.info(f"WS connected symbol={symbol}, total={len(self._all)}")
+
+    def disconnect(self, ws: WebSocket, symbol: str | None = None):
+        self._all.discard(ws)
+        if symbol:
+            self._subs[symbol].discard(ws)
+        else:
+            for s in list(self._subs.keys()):
+                self._subs[s].discard(ws)
+        logger.info(f"WS disconnected, total={len(self._all)}")
+
+    async def broadcast_quote(self, symbol: str, data: dict):
+        """Send quote update to all subscribers of a specific symbol."""
+        message = json.dumps({"type": "quote", "symbol": symbol, "data": data})
+        dead = set()
+        for ws in list(self._subs.get(symbol, [])):
+            try:
+                await ws.send_text(message)
+            except Exception:
+                dead.add(ws)
+        for ws in dead:
+            self.disconnect(ws, symbol)
+
+    async def broadcast_all(self, data: list[dict]):
+        """Broadcast heatmap snapshot to all connected clients."""
+        message = json.dumps({"type": "heatmap", "data": data})
+        dead = set()
+        for ws in list(self._all):
+            try:
+                await ws.send_text(message)
+            except Exception:
+                dead.add(ws)
+        for ws in dead:
+            self.disconnect(ws)
+
+
+manager = ConnectionManager()