97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
"""
|
|
WebSocket connection handler
|
|
Presentation Layer - manages WebSocket connections and real-time communication
|
|
"""
|
|
import asyncio
|
|
from typing import List
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class WebSocketManager:
|
|
"""Manages WebSocket connections and broadcasting"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: List[WebSocket] = []
|
|
|
|
async def connect(self, websocket: WebSocket) -> None:
|
|
"""Accept and store new WebSocket connection"""
|
|
await websocket.accept()
|
|
self.active_connections.append(websocket)
|
|
logger.info(f"New client connected. Total clients: {len(self.active_connections)}")
|
|
|
|
def disconnect(self, websocket: WebSocket) -> None:
|
|
"""Remove WebSocket connection"""
|
|
if websocket in self.active_connections:
|
|
self.active_connections.remove(websocket)
|
|
logger.info(f"Client disconnected. Total clients: {len(self.active_connections)}")
|
|
|
|
async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
|
|
"""Send message to specific WebSocket connection"""
|
|
try:
|
|
await websocket.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending personal message: {e}")
|
|
self.disconnect(websocket)
|
|
|
|
async def broadcast(self, message: str) -> None:
|
|
"""Broadcast message to all connected clients"""
|
|
if not self.active_connections:
|
|
return
|
|
|
|
try:
|
|
# Send to all connections concurrently
|
|
tasks = [
|
|
self._safe_send_message(connection, message)
|
|
for connection in self.active_connections.copy()
|
|
]
|
|
|
|
# Execute all sends concurrently and handle exceptions
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Remove failed connections
|
|
failed_connections = []
|
|
for i, result in enumerate(results):
|
|
if isinstance(result, Exception):
|
|
failed_connections.append(self.active_connections[i])
|
|
|
|
for connection in failed_connections:
|
|
self.disconnect(connection)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in broadcast: {e}")
|
|
|
|
async def _safe_send_message(self, websocket: WebSocket, message: str) -> None:
|
|
"""Safely send message to WebSocket with error handling"""
|
|
try:
|
|
await websocket.send_text(message)
|
|
except WebSocketDisconnect:
|
|
# Connection was closed
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error sending message to client: {e}")
|
|
raise
|
|
|
|
def get_connection_count(self) -> int:
|
|
"""Get number of active connections"""
|
|
return len(self.active_connections)
|
|
|
|
async def ping_all_connections(self) -> int:
|
|
"""Ping all connections to check health, return number of healthy connections"""
|
|
if not self.active_connections:
|
|
return 0
|
|
|
|
healthy_connections = []
|
|
for connection in self.active_connections.copy():
|
|
try:
|
|
await connection.ping()
|
|
healthy_connections.append(connection)
|
|
except Exception:
|
|
logger.debug("Removing unhealthy connection")
|
|
|
|
self.active_connections = healthy_connections
|
|
return len(healthy_connections)
|
|
|
|
# Global WebSocket manager instance
|
|
websocket_manager = WebSocketManager() |