39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
|
import json
|
||
|
|
import logging
|
||
|
|
from typing import Set
|
||
|
|
from fastapi import WebSocket
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class ConnectionManager:
|
||
|
|
def __init__(self):
|
||
|
|
self.active_connections: Set[WebSocket] = set()
|
||
|
|
|
||
|
|
async def connect(self, websocket: WebSocket):
|
||
|
|
await websocket.accept()
|
||
|
|
self.active_connections.add(websocket)
|
||
|
|
logger.info(f"WebSocket client connected. Total clients: {len(self.active_connections)}")
|
||
|
|
|
||
|
|
def disconnect(self, websocket: WebSocket):
|
||
|
|
self.active_connections.discard(websocket)
|
||
|
|
logger.info(f"WebSocket client disconnected. Remaining clients: {len(self.active_connections)}")
|
||
|
|
|
||
|
|
async def broadcast(self, message: dict):
|
||
|
|
if not self.active_connections:
|
||
|
|
return
|
||
|
|
payload = json.dumps(message, default=str)
|
||
|
|
stale_connections = set()
|
||
|
|
for connection in self.active_connections:
|
||
|
|
try:
|
||
|
|
await connection.send_text(payload)
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Failed to send to WebSocket client: {e}")
|
||
|
|
stale_connections.add(connection)
|
||
|
|
|
||
|
|
for stale in stale_connections:
|
||
|
|
self.active_connections.discard(stale)
|
||
|
|
|
||
|
|
|
||
|
|
manager = ConnectionManager()
|