gonets/SRV/api/v1/websocket.py

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)