k
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import json
|
|
import logging
|
|
from typing import List
|
|
|
|
from fastapi import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: List[WebSocket] = []
|
|
|
|
async def connect(self, ws: WebSocket):
|
|
await ws.accept()
|
|
self.active_connections.append(ws)
|
|
logger.info("WebSocket connected. Total connections: %d", len(self.active_connections))
|
|
|
|
async def disconnect(self, ws: WebSocket):
|
|
if ws in self.active_connections:
|
|
self.active_connections.remove(ws)
|
|
logger.info("WebSocket disconnected. Total connections: %d", len(self.active_connections))
|
|
|
|
async def broadcast(self, message: dict):
|
|
if not self.active_connections:
|
|
return
|
|
payload = json.dumps(message)
|
|
dead: List[WebSocket] = []
|
|
for connection in self.active_connections:
|
|
try:
|
|
await connection.send_text(payload)
|
|
except Exception as exc:
|
|
logger.warning("Failed to send to WebSocket: %s", exc)
|
|
dead.append(connection)
|
|
for ws in dead:
|
|
await self.disconnect(ws)
|
|
|
|
|
|
manager = ConnectionManager()
|