152 lines
No EOL
6 KiB
Python
152 lines
No EOL
6 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-серверу.
|
||
|
||
Обрабатывает подключение, отключение и трансляцию сообщений для клиентов,
|
||
подключённых через WebSocket.
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""Инициализирует менеджер подключений к WebSocket-серверу.
|
||
|
||
Создаёт пустое множество для хранения активных соединений.
|
||
"""
|
||
self.active_connections: Set[WebSocket] = set()
|
||
|
||
async def connect(self, websocket: WebSocket):
|
||
"""Подключает клиента к WebSocket-серверу.
|
||
|
||
Принимает подключение и добавляет WebSocket в список активных соединений.
|
||
|
||
Args:
|
||
websocket (WebSocket): Объект WebSocket для подключения.
|
||
|
||
Raises:
|
||
Exception: Если произошла ошибка при принятии подключения.
|
||
"""
|
||
await websocket.accept()
|
||
self.active_connections.add(websocket)
|
||
|
||
def disconnect(self, websocket: WebSocket):
|
||
"""Отключает клиента от WebSocket-серверу.
|
||
|
||
Удаляет WebSocket из списка активных соединений.
|
||
|
||
Args:
|
||
websocket (WebSocket): Объект WebSocket для отключения.
|
||
"""
|
||
self.active_connections.remove(websocket)
|
||
|
||
async def broadcast(self, message: dict):
|
||
"""Транслирует сообщения всем подключённым клиентам.
|
||
|
||
Отправляет JSON-данные всем активным соединениям, обрабатывая
|
||
отключения клиентов.
|
||
|
||
Args:
|
||
message (dict): Сообщение для трансляции в формате JSON.
|
||
|
||
Raises:
|
||
WebSocketDisconnect: Если клиент неожиданно отключился.
|
||
Exception: Если произошла другая ошибка при отправке.
|
||
"""
|
||
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-соединения.
|
||
|
||
Управляет соединением клиента, принимает данные для обновления позиций
|
||
контейнеров и транслирует сетевые данные.
|
||
|
||
Args:
|
||
websocket (WebSocket): Объект WebSocket для обработки соединения.
|
||
|
||
Raises:
|
||
WebSocketDisconnect: Если клиент отключился.
|
||
Exception: Если произошла ошибка при обработке данных.
|
||
"""
|
||
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():
|
||
"""Запуск таска для трансляции сетевых данных при старте приложения.
|
||
|
||
Создаёт асинхронную задачу для периодической трансляции данных.
|
||
|
||
Raises:
|
||
Exception: Если произошла ошибка при запуске таска.
|
||
"""
|
||
asyncio.create_task(broadcast_data())
|
||
|
||
|
||
async def broadcast_data():
|
||
"""Таск для периодической трансляции сетевых данных.
|
||
|
||
Выполняет сканирование сети и отправляет данные всем подключённым клиентам
|
||
с интервалом в 2 секунды.
|
||
|
||
Raises:
|
||
Exception: Если произошла ошибка при получении данных или трансляции.
|
||
"""
|
||
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) |