import asyncio import logging from typing import Set from fastapi import APIRouter, WebSocket, WebSocketDisconnect from starlette.websockets import WebSocketState from .network import update_container_position, get_network_data router = APIRouter() logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class WebSocketManager: """ Класс для управления подключениями к WebSocket-серверу """ def __init__(self): """Инициализация менеджера подключений к WebSocket-серверу""" self.active_connections: Set[WebSocket] = set() async def connect(self, websocket: WebSocket): """Подключение клиента к WebSocket-серверу""" await websocket.accept() self.active_connections.add(websocket) def disconnect(self, websocket: WebSocket): """Отключение клиента от WebSocket-серверу""" self.active_connections.remove(websocket) async def broadcast(self, message: dict): """Трансляция сообщений всем подключенным клиентам""" for connection in list(self.active_connections): try: if connection.client_state == WebSocketState.CONNECTED: await connection.send_json(message) except WebSocketDisconnect: self.disconnect(connection) except Exception as e: logger.error(f"Сообщение об ошибке трансляции: {e}") self.disconnect(connection) manager = WebSocketManager() @router.websocket("/api/v1") async def websocket_endpoint(websocket: WebSocket): """Обработчик WebSocket-соединения""" await manager.connect(websocket) try: while True: data = await websocket.receive_json() logger.info(f"Полученные данные: {data}") if all(key in data for key in ["container_id", "container_x", "container_y"]): await update_container_position(data) network_data = await get_network_data() logger.info(f"Текущие сетевые данные: {network_data}") await manager.broadcast(network_data) await asyncio.sleep(2) except WebSocketDisconnect: logger.info("Клиент отключен") except Exception as e: logger.error(f"Ошибка в websocket_endpoint:{e}") finally: manager.disconnect(websocket) if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close() @router.on_event("startup") async def startup_event(): """Запуск таска для трансляции сетевых данных""" asyncio.create_task(broadcast_data()) async def broadcast_data(): """Таск для трансляции сетевых данных""" while True: try: network_data = await get_network_data() if "activation" not in network_data: logger.error( "Ошибка: Ключ 'activation' отсутствует в полученных данных сети" ) await asyncio.sleep(2) continue await manager.broadcast(network_data) except Exception as e: logger.error(f"Ошибка в broadcast_data:{e}") await asyncio.sleep(2)