Загрузка данных


вс кли 
import asyncio
import logging
import time
from typing import Any, Callable, List, Optional

import msgpack
import requests
import websockets
from websockets import InvalidStatus

from constants.architecture_constants import WebSocketClientConstants as WS_Const
from utils.msgpack_utils.message_filters import is_desired_invocation_id, is_desired_type
from utils.msgpack_utils.msgpack_utils import encode_with_varint_prefix, parse_message

logger = logging.getLogger(__name__)


class WebSocketClient:
    """
    Асинхронный ws-клиент для api-gateway по протоколу Async Api
    """

    def __init__(
        self,
        host: str,
        access_token: str,
        x_user_id: str,
        reconnect_interval: float = WS_Const.DEFAULT_RECONNECT_INTERVAL,
    ):
        self._host = host
        self._access_token = access_token
        self._x_user_id = x_user_id
        self._reconnect_interval = reconnect_interval
        self._ws_url = f"wss://{host.rstrip('/')}{WS_Const.WS_HUBS}"
        self._negotiate_url = f"https://{host.rstrip('/')}{WS_Const.NEGOTIATE_ENDPOINT}"
        self._buffer = b""
        self._next_id = WS_Const.START_INVOCATION_ID

        self._ws: websockets.ClientConnection | None = None
        self.recv_queue: asyncio.Queue[Any] = asyncio.Queue()
        self._recv_task: asyncio.Task | None = None
        self._ping_task: asyncio.Task | None = None
        self._stop_event = asyncio.Event()
        self._invocation_id: Optional[str] = None
        self.suppress_recv_logging: bool = False
        self.received_messages_count: int = 0
        self.sent_ping_count: int = 0

    @property
    def invocation_id(self):
        return self._invocation_id

    @property
    def is_connected(self) -> bool:
        """
        True, если wss-соединение установлено и приём сообщений работает.

        Соединение считается живым, пока открыт сокет и фоновый приёмник
        (_recv_task) не завершился. При обрыве связи _recv_loop возвращается,
        и задача переходит в состояние done - это признак потери соединения.
        """
        if self._ws is None:
            return False
        if self._recv_task is None:
            return False
        return not self._recv_task.done()

    def clear_queue(self):
        """
        Очищает очередь путем пересоздания экземпляра класса очереди
        """
        while not self.recv_queue.empty():
            try:
                self.recv_queue.get_nowait()
                self.recv_queue.task_done()
            except asyncio.QueueEmpty as message_empty:
                logger.info(f"Очередь сообщений очищена: {message_empty}")
                break

    def update_credentials(self, access_token: str, x_user_id: str) -> None:
        """
        Обновляет креды для последующего reconnect/negotiate.

        Нужно при смене датасета: токен и x_user_id обновляются в group_state,
        а соединение переиспользуется (пересоздаётся через reconnect).
        """
        self._access_token = access_token
        self._x_user_id = x_user_id

    async def __aenter__(self):
        await self._connect_loop()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        self._stop_event.set()
        if self._ping_task:
            self._ping_task.cancel()
            try:
                await self._ping_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()
        if self._recv_task:
            await self._recv_task

        # TODO: дергать ручку завершения сессии в LDS-4083

    async def _handshake(self) -> None:
        payload = WS_Const.HANDSHAKE_MESSAGE + WS_Const.RS.decode()
        logger.debug(
            f"Отправлен handshake: {payload}",
        )
        await self._ws.send(payload)

        buf = self._buffer
        finish = time.monotonic() + WS_Const.HANDSHAKE_WAITING
        while time.monotonic() < finish:
            chunk = await self._ws.recv()
            logger.info(f"Ответ на handshake: {chunk}")
            buf += chunk.encode() if isinstance(chunk, str) else chunk
            if WS_Const.RS in buf:
                return

        raise TimeoutError("Handshake timeout: не получили сообщение с разделителем RS за указанное время")

    async def negotiate(self) -> str:
        """
        SignalR-negotiate: получает connectionToken для ws-подключения.

        connectionToken подставляется в параметр id= URL вебсокета. Без negotiate
        api-gateway не регистрирует соединение как полноценного SignalR-клиента,
        поэтому он обязателен для корректного подключения.
        """
        params = {
            "token": self._access_token,
            "xUserId": self._x_user_id,
            "negotiateVersion": WS_Const.NEGOTIATE_VERSION,
        }
        response = await asyncio.to_thread(
            requests.post,
            self._negotiate_url,
            params=params,
            timeout=WS_Const.NEGOTIATE_TIMEOUT_SECONDS,
        )
        response.raise_for_status()
        data = response.json()
        connection_token = data.get("connectionToken")
        logger.info(
            "[NEGOTIATE] POST %s -> %s (connectionId=%s, connectionToken=%s)",
            WS_Const.NEGOTIATE_ENDPOINT,
            response.status_code,
            data.get("connectionId"),
            connection_token,
        )
        if not connection_token:
            raise RuntimeError("negotiate не вернул connectionToken")
        return connection_token

    async def _connect_loop(self, timeout: float = WS_Const.WS_CONNECT_TIMEOUT_SECONDS) -> None:
        """
        Цикл подключения с повторными попытками до наступления stop_event
        или истечения WS_CONNECT_TIMEOUT_SECONDS.
        """
        deadline = time.monotonic() + timeout
        attempt = 0
        transient_errors = (ConnectionError, OSError, asyncio.TimeoutError, InvalidStatus)

        while not self._stop_event.is_set():
            attempt += 1
            try:
                connection_token = await self.negotiate()
                self.ws_request = (
                    f"{self._ws_url}/?token={self._access_token}"
                    f"&xUserId={self._x_user_id}&id={connection_token}"
                )
                logger.info(f"Попытка подключения по wss: {self._ws_url}/?token=...&xUserId=...&id=...")
                self._ws = await websockets.connect(
                    self.ws_request,
                    ping_interval=WS_Const.PING_INTERVAL,
                    ping_timeout=WS_Const.PING_TIMEOUT,
                    close_timeout=WS_Const.CLOSE_TIMEOUT,
                )
                # Handshake
                await self._handshake()
                # Запускаем приём и SignalR keepalive в фоне
                self._recv_task = asyncio.create_task(self._recv_loop())
                self._ping_task = asyncio.create_task(self._ping_loop())
                logger.info("Websocket connected")
                return
            except transient_errors as exc:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError(
                        f"WSS hub не готов за {WS_Const.WS_CONNECT_TIMEOUT_SECONDS} с " f"(попыток: {attempt}): {exc}"
                    ) from exc
                status_info = ""
                if isinstance(exc, InvalidStatus):
                    response = getattr(exc, "response", None)
                    if response is not None:
                        status_info = f", HTTP {response.status_code}"
                logger.warning(
                    "WSS подключение не установлено (попытка %s%s): %s. " "Повтор через %s с, осталось %.0f с",
                    attempt,
                    status_info,
                    exc,
                    self._reconnect_interval,
                    remaining,
                )
                await asyncio.sleep(self._reconnect_interval)

    async def _recv_loop(self) -> None:
        """
        Прием сообщений, парсинг и отправка в очередь.
        """
        assert self._ws is not None

        while not self._stop_event.is_set():
            try:
                chunk = await self._ws.recv()
                if not self.suppress_recv_logging:
                    logger.debug(f"Сырые биты до обработки: {chunk[:100]}")
            except websockets.ConnectionClosed as e:
                logger.warning(f"WebSocket соединение разорвано: {e}")
                return

            result_message = parse_message(chunk)
            await self.recv_queue.put(result_message)
            self.received_messages_count += 1
            if not self.suppress_recv_logging:
                str_message = str(result_message)
                logger.info(
                    f"Обработанное сообщение от api-gateway: {str_message[:200]}... полное сообщение в attach",
                )
            continue

    async def send_ping(self) -> None:
        """Отправляет SignalR Ping [6] для прикладного keepalive."""
        if not self._ws:
            return
        packet = encode_with_varint_prefix(msgpack.packb([WS_Const.PING_MESSAGE_TYPE], use_bin_type=True))
        await self._ws.send(packet)
        self.sent_ping_count += 1

    async def _ping_loop(self) -> None:
        """Периодически шлёт SignalR Ping, пока соединение живо."""
        while not self._stop_event.is_set():
            await asyncio.sleep(WS_Const.SIGNALR_PING_INTERVAL_SECONDS)
            if self._stop_event.is_set() or self._ws is None:
                return
            try:
                await self.send_ping()
            except (websockets.ConnectionClosed, OSError):
                return

    async def invoke(self, target: str, args: list) -> None:
        """
        Отправляет удаленный вызов invocation о websocket соединению

        Метод формирует сообщение по протоколу SignalR, включая:
        - типа сообщения
        - заголовки
        - уникальный идентификатор запроса
        - имя целевого метода
        - аргументы запроса

        Сообщение запаковывается в messagepack и отправляется через текущее websocket соединение
        """

        if not self._ws:
            raise websockets.WebSocketException("Не установлено подключение по wss")
        self._invocation_id = str(self._next_id)
        self._next_id += 1
        invocation = [
            WS_Const.DEFAULT_SIGNALR_MESSAGE_TYPE,
            WS_Const.DEFAULT_SIGNALR_MAP_HEADERS,
            self._invocation_id,
            target,
            [args],
        ]
        logger.info(f"Сообщение подготовлено к отправке: {invocation}")
        payload = msgpack.packb(invocation, use_bin_type=True)
        packet = encode_with_varint_prefix(payload)
        logger.debug(f"Отправляем сообщение: {packet}")
        await self._ws.send(packet)

    async def invoke_stream(self, target: str, args: list) -> None:
        """
        Отправляет streaming-вызов (StreamInvocation) по протоколу SignalR.
        """
        if not self._ws:
            raise websockets.WebSocketException("Не установлено подключение по wss")
        self._invocation_id = str(self._next_id)
        self._next_id += 1
        invocation = [
            WS_Const.STREAM_INVOCATION_MESSAGE_TYPE,
            WS_Const.DEFAULT_SIGNALR_MAP_HEADERS,
            self._invocation_id,
            target,
            [args],
        ]
        logger.info(f"Streaming-сообщение подготовлено к отправке: {invocation}")
        payload = msgpack.packb(invocation, use_bin_type=True)
        packet = encode_with_varint_prefix(payload)
        await self._ws.send(packet)

    async def receive_by_type(self, message_type: str, timeout: float = WS_Const.FILTERING_TIMEOUT) -> List[Any]:
        """
        Фильтрует сообщения по message_type
        """
        try:
            return await self._receive_by(filter_func=lambda msg: is_desired_type(msg, message_type), timeout=timeout)
        except websockets.WebSocketException:
            raise websockets.WebSocketException(f"Ошибка при фильтрации сообщений по {message_type}")

    async def receive_by_invocation_id(
        self, invocation_id: str, timeout: float = WS_Const.FILTERING_TIMEOUT
    ) -> List[Any]:
        """
        Фильтрует сообщения по invocation_id
        """
        try:
            return await self._receive_by(
                filter_func=lambda msg: is_desired_invocation_id(msg, invocation_id), timeout=timeout
            )
        except websockets.WebSocketException:
            raise websockets.WebSocketException("Ошибка при фильтрации сообщений по invocation_id")

    async def _receive_by(self, filter_func: Callable[[list], bool], timeout: float) -> List[Any]:
        """
        Ждет и фильтрует сообщение по filter_func
        """
        # 1) Единая точка вычисления дедлайна
        deadline = time.monotonic() + timeout

        while True:
            # 2) Остаток времени до таймаута
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise asyncio.TimeoutError(f"Timeout при фильтрации сообщений {timeout:.1f} секунд")

            try:
                # 3) Получает сообщение
                msg = await asyncio.wait_for(self.recv_queue.get(), timeout=remaining)
            except asyncio.TimeoutError:
                # 4) Явно перехватывает и пробрасывает свой TimeoutError
                raise asyncio.TimeoutError(f"Timeout при фильтрации сообщений {timeout:.1f} секунд")

            # 5) Фильтрация по filter_func
            if isinstance(msg, list) and filter_func(msg):
                return msg

    async def reconnect(self):
        """
        Переподключение websocket соединения
        """
        try:
            if self._ping_task:
                self._ping_task.cancel()
                try:
                    await self._ping_task
                except asyncio.CancelledError:
                    pass
                self._ping_task = None

            if self._recv_task:
                try:
                    self._recv_task.cancel()
                    await self._recv_task
                except asyncio.CancelledError:
                    pass
                self._recv_task = None

            if self._ws:
                await self._ws.close()
                self._ws = None
            await self._connect_loop(WS_Const.WS_RECONNECT_TIMEOUT_SECONDS)
            return self
        except (asyncio.TimeoutError, ConnectionError, ConnectionResetError, OSError) as error:
            raise RuntimeError(f"Ошибка при попытке переподключения {error}")






















арх конст
    WS_SESSION_SUPERVISOR_INTERVAL_SECONDS: float = 1.0  # периодичность проверки жизни соединения
    WS_SESSION_HEARTBEAT_INTERVAL_SECONDS: float = 60.0  # периодичность heartbeat-лога сессии





















вс сешн
"""
Постоянная ws-сессия на датасет, имитирующая подключение фронта.
"""

from __future__ import annotations

import asyncio
import logging
import uuid

import requests

from clients.websocket_client import WebSocketClient
from constants.architecture_constants import HTTPClientConstants as HttpConst
from constants.architecture_constants import StandConstants
from constants.architecture_constants import WebSocketClientConstants as WS_Const

logger = logging.getLogger(__name__)


class WsSessionManager:
    """Управляет постоянным ws-соединением датасета: подписки, keepalive, реконнект."""

    def __init__(self, group_state: dict, client: WebSocketClient) -> None:
        self._group_state = group_state
        self._client = client
        self._stop_event = asyncio.Event()
        self._generation = group_state.get("dataset_generation")
        self._correlation_id = str(uuid.uuid4())
        self._supervisor_task: asyncio.Task | None = None
        self._rest_ping_task: asyncio.Task | None = None
        self._heartbeat_task: asyncio.Task | None = None
        self._rest_ping_count = 0
        self._last_received = 0
        self._last_sent_ping = 0
        self._last_rest_ping = 0

    async def start(self) -> None:
        """Подписывается на bundle и запускает фоновые задачи."""
        await self._subscribe_bundle()
        self._supervisor_task = asyncio.create_task(self._supervisor_loop())
        self._rest_ping_task = asyncio.create_task(self._rest_ping_loop())
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        logger.info("[WS_SESSION] сессия запущена (correlation_id=%s)", self._correlation_id)

    async def stop(self) -> None:
        """Останавливает фоновые задачи."""
        self._stop_event.set()
        for task in (self._rest_ping_task, self._supervisor_task, self._heartbeat_task):
            if task is not None:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        logger.info("[WS_SESSION] сессия остановлена")

    def _current_tu_id(self) -> int | None:
        """tuId для MainPage: resolved из Администрирования, либо legacy из маркера."""
        return self._group_state.get("resolved_tu_id") or self._group_state.get("legacy_tu_id")

    async def _subscribe_bundle(self) -> None:
        """Подписывается на bundle так же, как это делает фронт."""
        tu_id = self._current_tu_id()
        subscriptions = [
            (WS_Const.SUBSCRIBE_PLAYER_STATUS_REQUEST, []),
            (WS_Const.SUBSCRIBE_SELF_DIAGNOSTIC_REQUEST, []),
        ]
        if tu_id is not None:
            subscriptions.append(
                (WS_Const.SUBSCRIBE_MAIN_PAGE_INFO_REQUEST, {"tuIds": [tu_id], "additionalProperties": None})
            )
        for invoke_type, invoke_params in subscriptions:
            try:
                await self._client.invoke(invoke_type, invoke_params)
            except Exception as error:
                logger.warning("[WS_SESSION] не удалось подписаться на %s: %s", invoke_type, error)
        logger.info("[WS_SESSION] bundle-подписки оформлены (tu_id=%s)", tu_id)

    async def _supervisor_loop(self) -> None:
        """Следит за состоянием соединения и сменой датасета."""
        while not self._stop_event.is_set():
            await asyncio.sleep(WS_Const.WS_SESSION_SUPERVISOR_INTERVAL_SECONDS)
            if self._stop_event.is_set():
                return
            generation = self._group_state.get("dataset_generation")
            if not self._client.is_connected or generation != self._generation:
                logger.info(
                    "[WS_SESSION] переподключение: is_connected=%s, generation %s -> %s",
                    self._client.is_connected,
                    self._generation,
                    generation,
                )
                await self._reestablish(generation)

    async def _reestablish(self, new_generation: int | None) -> None:
        """Переподключается с актуальными кредами и переподписывается на bundle."""
        auth_token = self._group_state.get("auth_token")
        x_user_id = self._group_state.get("x_user_id")
        if not auth_token or not x_user_id:
            logger.warning("[WS_SESSION] нет кредов для переподключения, пропускаю")
            return
        self._client.update_credentials(auth_token, x_user_id)
        try:
            await self._client.reconnect()
        except Exception as error:
            logger.warning("[WS_SESSION] не удалось переподключиться: %s", error)
            return
        self._generation = new_generation
        await self._subscribe_bundle()
        logger.info("[WS_SESSION] соединение переподключено, bundle переподписан")

    async def _heartbeat_loop(self) -> None:
        """Раз в минуту логирует состояние сессии: соединение, очередь, активность."""
        while not self._stop_event.is_set():
            await asyncio.sleep(WS_Const.WS_SESSION_HEARTBEAT_INTERVAL_SECONDS)
            if self._stop_event.is_set():
                return
            received = self._client.received_messages_count
            sent_ping = self._client.sent_ping_count
            rest_ping = self._rest_ping_count
            logger.info(
                "[WS_SESSION] heartbeat: активно=%s, очередь=%d, получено_за_минуту=%d (всего=%d), "
                "signalr_ping=%d, rest_ping=%d",
                self._client.is_connected,
                self._client.recv_queue.qsize(),
                received - self._last_received,
                received,
                sent_ping - self._last_sent_ping,
                rest_ping - self._last_rest_ping,
            )
            self._last_received = received
            self._last_sent_ping = sent_ping
            self._last_rest_ping = rest_ping

    async def _rest_ping_loop(self) -> None:
        """Периодически пингует REST /Ping с correlation-id."""
        while not self._stop_event.is_set():
            await asyncio.sleep(WS_Const.REST_PING_INTERVAL_SECONDS)
            if self._stop_event.is_set():
                return
            await self._ping_rest()

    async def _ping_rest(self) -> None:
        """Один REST-пинг. Ошибки (5xx/4xx/сеть) не роняют keepalive."""
        stand_host = self._group_state.get("stand_host")
        auth_token = self._group_state.get("auth_token")
        if not stand_host or not auth_token:
            logger.warning("[REST_PING] нет stand_host/auth_token, пропускаю")
            return
        url = f"https://{stand_host}{StandConstants.PING_URL_PATH}"
        headers = {
            HttpConst.X_SECURITY_SIGNATURE_KEY: auth_token,
            HttpConst.X_TN_CORRELATION_ID_KEY: self._correlation_id,
        }
        self._rest_ping_count += 1
        try:
            response = await asyncio.to_thread(requests.post, url, headers=headers, timeout=10)
            if not response.ok:
                logger.warning(
                    "[REST_PING] %s -> %s: %s",
                    StandConstants.PING_URL_PATH,
                    response.status_code,
                    response.text[:200],
                )
            else:
                logger.info("[REST_PING] %s -> %s", StandConstants.PING_URL_PATH, response.status_code)
        except requests.RequestException as error:
            logger.warning("[REST_PING] ошибка запроса %s: %s", StandConstants.PING_URL_PATH, error)












пайпроджект
asyncio_mode="auto"
asyncio_default_fixture_loop_scope="session"
asyncio_default_test_loop_scope="session"