2025-12-12 13:30:09 +03:00
|
|
|
|
from typing import Dict, Set
|
|
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
from uuid import UUID
|
|
|
|
|
|
import json
|
2025-12-19 20:46:00 +03:00
|
|
|
|
import time
|
|
|
|
|
|
import asyncio
|
2025-12-12 13:30:09 +03:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ConnectionManager:
|
|
|
|
|
|
def __init__(self):
|
|
|
|
|
|
# room_id -> set of (websocket, user_id)
|
|
|
|
|
|
self.active_connections: Dict[UUID, Set[tuple[WebSocket, UUID]]] = {}
|
2025-12-19 20:46:00 +03:00
|
|
|
|
# (room_id, user_id) -> last pong timestamp
|
|
|
|
|
|
self.last_pong: Dict[tuple[UUID, UUID], float] = {}
|
|
|
|
|
|
# Ping interval and timeout in seconds
|
|
|
|
|
|
self.ping_interval = 30 # Отправлять ping каждые 30 секунд
|
|
|
|
|
|
self.ping_timeout = 60 # Отключать если нет ответа 60 секунд
|
2025-12-12 13:30:09 +03:00
|
|
|
|
|
|
|
|
|
|
async def connect(self, websocket: WebSocket, room_id: UUID, user_id: UUID):
|
|
|
|
|
|
await websocket.accept()
|
|
|
|
|
|
if room_id not in self.active_connections:
|
|
|
|
|
|
self.active_connections[room_id] = set()
|
|
|
|
|
|
self.active_connections[room_id].add((websocket, user_id))
|
2025-12-19 20:46:00 +03:00
|
|
|
|
# Инициализируем время последнего pong
|
|
|
|
|
|
self.last_pong[(room_id, user_id)] = time.time()
|
2025-12-12 13:30:09 +03:00
|
|
|
|
|
|
|
|
|
|
def disconnect(self, websocket: WebSocket, room_id: UUID, user_id: UUID):
|
|
|
|
|
|
if room_id in self.active_connections:
|
|
|
|
|
|
self.active_connections[room_id].discard((websocket, user_id))
|
|
|
|
|
|
if not self.active_connections[room_id]:
|
|
|
|
|
|
del self.active_connections[room_id]
|
2025-12-19 20:46:00 +03:00
|
|
|
|
# Удаляем запись о последнем pong
|
|
|
|
|
|
self.last_pong.pop((room_id, user_id), None)
|
|
|
|
|
|
|
|
|
|
|
|
def update_pong(self, room_id: UUID, user_id: UUID):
|
|
|
|
|
|
"""Обновить время последнего pong от пользователя"""
|
|
|
|
|
|
self.last_pong[(room_id, user_id)] = time.time()
|
2025-12-12 13:30:09 +03:00
|
|
|
|
|
|
|
|
|
|
async def broadcast_to_room(self, room_id: UUID, message: dict, exclude_user: UUID = None):
|
|
|
|
|
|
if room_id not in self.active_connections:
|
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
message_json = json.dumps(message, default=str)
|
|
|
|
|
|
disconnected = []
|
|
|
|
|
|
|
|
|
|
|
|
for websocket, user_id in self.active_connections[room_id]:
|
|
|
|
|
|
if exclude_user and user_id == exclude_user:
|
|
|
|
|
|
continue
|
|
|
|
|
|
try:
|
|
|
|
|
|
await websocket.send_text(message_json)
|
|
|
|
|
|
except Exception:
|
|
|
|
|
|
disconnected.append((websocket, user_id))
|
|
|
|
|
|
|
|
|
|
|
|
for conn in disconnected:
|
|
|
|
|
|
self.active_connections[room_id].discard(conn)
|
|
|
|
|
|
|
2025-12-19 20:46:00 +03:00
|
|
|
|
async def send_to_user(self, room_id: UUID, user_id: UUID, message: dict):
|
|
|
|
|
|
"""Отправить сообщение конкретному пользователю в комнате"""
|
|
|
|
|
|
if room_id not in self.active_connections:
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
message_json = json.dumps(message, default=str)
|
|
|
|
|
|
for websocket, uid in self.active_connections[room_id]:
|
|
|
|
|
|
if uid == user_id:
|
|
|
|
|
|
try:
|
|
|
|
|
|
await websocket.send_text(message_json)
|
|
|
|
|
|
return True
|
|
|
|
|
|
except Exception:
|
|
|
|
|
|
return False
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
2025-12-12 13:30:09 +03:00
|
|
|
|
def get_room_user_count(self, room_id: UUID) -> int:
|
|
|
|
|
|
if room_id not in self.active_connections:
|
|
|
|
|
|
return 0
|
|
|
|
|
|
return len(self.active_connections[room_id])
|
|
|
|
|
|
|
2025-12-19 20:46:00 +03:00
|
|
|
|
def get_room_users(self, room_id: UUID) -> Set[UUID]:
|
|
|
|
|
|
"""Получить список user_id в комнате"""
|
|
|
|
|
|
if room_id not in self.active_connections:
|
|
|
|
|
|
return set()
|
|
|
|
|
|
return {user_id for _, user_id in self.active_connections[room_id]}
|
|
|
|
|
|
|
|
|
|
|
|
def get_all_connections(self) -> Dict[UUID, Set[tuple[WebSocket, UUID]]]:
|
|
|
|
|
|
"""Получить все активные соединения (для фоновой задачи пинга)"""
|
|
|
|
|
|
return self.active_connections
|
|
|
|
|
|
|
2025-12-12 13:30:09 +03:00
|
|
|
|
|
|
|
|
|
|
manager = ConnectionManager()
|