crypto_monitor/lib/exchange/okx_public_ws_lib.py

"""OKX 公共 WebSocket(同步线程):订阅 tickers / index-tickers,自动重连."""
from __future__ import annotations

import json
import logging
import threading
import time
from collections.abc import Callable
from typing import Any

logger = logging.getLogger(__name__)

OKX_PUBLIC_WS_URL = "wss://ws.okx.com:8443/ws/v5/public"
_SUBSCRIBE_CHUNK = 40
_APP_PING_SEC = 20.0


class OkxPublicWs:
    """单连接公共 WS;set_subscriptions 全量对齐目标频道."""

    def __init__(
        self,
        *,
        on_data: Callable[[dict[str, Any]], None],
        url: str = OKX_PUBLIC_WS_URL,
        name: str = "okx-public-ws",
    ) -> None:
        self._on_data = on_data
        self._url = url
        self._name = name
        self._lock = threading.RLock()
        self._desired: dict[str, dict[str, str]] = {}
        self._active: set[str] = set()
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._ws: Any = None
        self._connected = False
        self._last_msg_at = 0.0

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def last_msg_at(self) -> float:
        return self._last_msg_at

    def start(self) -> None:
        if self._thread and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run_loop, name=self._name, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        ws = self._ws
        if ws is not None:
            try:
                ws.close()
            except Exception:
                pass
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=3.0)

    def set_subscriptions(self, args: list[dict[str, str]]) -> None:
        desired: dict[str, dict[str, str]] = {}
        for raw in args:
            if not isinstance(raw, dict):
                continue
            channel = str(raw.get("channel") or "").strip()
            inst_id = str(raw.get("instId") or "").strip()
            if not channel or not inst_id:
                continue
            key = f"{channel}:{inst_id}"
            desired[key] = {"channel": channel, "instId": inst_id}
        with self._lock:
            self._desired = desired
            ws = self._ws
            connected = self._connected
            active = set(self._active)
        if connected and ws is not None:
            self._sync_subs(ws, active, desired)

    def _sync_subs(
        self,
        ws: Any,
        active: set[str],
        desired: dict[str, dict[str, str]],
    ) -> None:
        unsub_args: list[dict[str, str]] = []
        for key in active - set(desired.keys()):
            channel, _, inst_id = key.partition(":")
            if channel and inst_id:
                unsub_args.append({"channel": channel, "instId": inst_id})
        sub_args = [desired[k] for k in (set(desired.keys()) - active)]
        if unsub_args:
            self._send_op(ws, "unsubscribe", unsub_args)
        if sub_args:
            self._send_op(ws, "subscribe", sub_args)
        with self._lock:
            self._active = set(desired.keys())

    def _send_op(self, ws: Any, op: str, args: list[dict[str, str]]) -> None:
        for i in range(0, len(args), _SUBSCRIBE_CHUNK):
            chunk = args[i : i + _SUBSCRIBE_CHUNK]
            try:
                ws.send(json.dumps({"op": op, "args": chunk}, ensure_ascii=False))
            except Exception as e:
                logger.warning("%s %s failed: %s", self._name, op, e)
                return
            if i + _SUBSCRIBE_CHUNK < len(args):
                time.sleep(0.08)

    def _run_loop(self) -> None:
        try:
            import websocket
        except ImportError:
            logger.error("%s: websocket-client not installed", self._name)
            return
        backoff = 1.0
        while not self._stop.is_set():
            opened = False
            try:
                self._connected = False
                with self._lock:
                    self._active.clear()

                def on_open(ws: Any) -> None:
                    nonlocal opened
                    opened = True
                    self._connected = True
                    self._last_msg_at = time.time()
                    with self._lock:
                        desired = dict(self._desired)
                    self._sync_subs(ws, set(), desired)

                def on_message(_ws: Any, message: str) -> None:
                    self._last_msg_at = time.time()
                    if message == "pong":
                        return
                    try:
                        payload = json.loads(message)
                    except Exception:
                        return
                    if not isinstance(payload, dict):
                        return
                    if payload.get("event") in ("subscribe", "unsubscribe", "error"):
                        if payload.get("event") == "error":
                            logger.warning("%s event error: %s", self._name, payload)
                        return
                    if payload.get("arg") and payload.get("data") is not None:
                        try:
                            self._on_data(payload)
                        except Exception:
                            logger.exception("%s on_data failed", self._name)

                def on_error(_ws: Any, error: Any) -> None:
                    logger.warning("%s error: %s", self._name, error)

                def on_close(_ws: Any, *_args: Any) -> None:
                    self._connected = False

                self._ws = websocket.WebSocketApp(
                    self._url,
                    on_open=on_open,
                    on_message=on_message,
                    on_error=on_error,
                    on_close=on_close,
                )
                ping_stop = threading.Event()

                def ping_loop() -> None:
                    while not self._stop.is_set() and not ping_stop.is_set():
                        ws = self._ws
                        if ws is not None and self._connected:
                            try:
                                ws.send("ping")
                            except Exception:
                                pass
                        if ping_stop.wait(_APP_PING_SEC):
                            break

                ping_thread = threading.Thread(
                    target=ping_loop, name=f"{self._name}-ping", daemon=True
                )
                ping_thread.start()
                self._ws.run_forever(ping_interval=0)
                ping_stop.set()
            except Exception as e:
                logger.warning("%s run failed: %s", self._name, e)
            finally:
                self._connected = False
                self._ws = None
            if self._stop.is_set():
                break
            time.sleep(backoff)
            backoff = 1.0 if opened else min(30.0, backoff * 1.7)