90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
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)
|