crypto_okx_hub/lib/overview_cache_lib.py

"""中控总览热缓存:后台常驻并行 pull,与浏览器是否打开无关;SSE 仅扇出版本通知."""
from __future__ import annotations

import json
import os
import queue
import threading
import time
from collections.abc import Callable, Iterator
from typing import Any

OnPullFn = Callable[[list[dict[str, Any]]], None]


def _poll_interval_sec() -> float:
    try:
        return max(3.0, float(os.getenv("HUB_PULL_INTERVAL_SEC") or 10))
    except (TypeError, ValueError):
        return 10.0


def _sse_heartbeat_sec() -> float:
    try:
        return max(10.0, float(os.getenv("HUB_SSE_HEARTBEAT_SEC") or 25))
    except (TypeError, ValueError):
        return 25.0


class OverviewStore:
    def __init__(self) -> None:
        self._lock = threading.RLock()
        self.version = 0
        self.instances: list[dict[str, Any]] = []
        self.pulling = True  # 启动至首轮完成前视为拉取中
        self.last_error: str | None = None
        self.updated_at: str = ""
        self._ready = False
        self._subscribers: list[queue.Queue[str | None]] = []
        self._stop = threading.Event()
        self._refresh = threading.Event()
        self._thread: threading.Thread | None = None
        self._pull_fn: Callable[[], list[dict[str, Any]]] | None = None
        self._on_pull: OnPullFn | None = None

    def start(
        self,
        pull_fn: Callable[[], list[dict[str, Any]]],
        *,
        on_pull: OnPullFn | None = None,
    ) -> None:
        self._pull_fn = pull_fn
        self._on_pull = on_pull
        if self._thread and self._thread.is_alive():
            return
        self._stop.clear()
        with self._lock:
            self.pulling = True
            self._ready = False
        self._thread = threading.Thread(
            target=self._loop,
            daemon=True,
            name="hub-overview-pull",
        )
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        self._refresh.set()
        self._broadcast(close=True)

    def request_refresh(self) -> None:
        self._refresh.set()

    def snapshot_dict(self) -> dict[str, Any]:
        with self._lock:
            rows = list(self.instances)
            ver = self.version
            pulling = self.pulling
            err = self.last_error
            updated_at = self.updated_at
        return {
            "ok": True,
            "instances": rows,
            "overview_version": ver,
            "pulling": pulling,
            "updated_at": updated_at,
            "error": err,
            "poll_interval_sec": _poll_interval_sec(),
            "from_cache": True,
        }

    def event_dict(self) -> dict[str, Any]:
        with self._lock:
            return {
                "overview_version": self.version,
                "updated_at": self.updated_at,
                "pulling": self.pulling,
                "ok": True,
                "error": self.last_error,
                "poll_interval_sec": _poll_interval_sec(),
            }

    def _loop(self) -> None:
        assert self._pull_fn is not None
        while not self._stop.is_set():
            self._pull_once()
            if self._stop.is_set():
                break
            self._refresh.clear()
            self._refresh.wait(timeout=_poll_interval_sec())

    def _pull_once(self) -> None:
        assert self._pull_fn is not None
        with self._lock:
            self.pulling = True
        self._broadcast()
        try:
            rows = self._pull_fn()
            if not isinstance(rows, list):
                rows = []
            if self._on_pull is not None:
                try:
                    self._on_pull(rows)
                except Exception:
                    pass
            with self._lock:
                self.version += 1
                self.instances = rows
                self.updated_at = time.strftime("%Y-%m-%d %H:%M:%S")
                self.last_error = None
                self.pulling = False
                self._ready = True
        except Exception as e:
            with self._lock:
                self.version += 1
                self.last_error = str(e) or "pull_failed"
                self.pulling = False
                self._ready = True
                # 保留上一轮成功缓存,避免整页空白
        self._broadcast()

    def _broadcast(self, *, close: bool = False) -> None:
        with self._lock:
            subs = list(self._subscribers)
            event = None if close else json.dumps(self.event_dict(), ensure_ascii=False)
        dead: list[queue.Queue[str | None]] = []
        for q in subs:
            try:
                q.put_nowait(None if close else event)
            except queue.Full:
                try:
                    q.get_nowait()
                except queue.Empty:
                    pass
                try:
                    q.put_nowait(event)
                except queue.Full:
                    dead.append(q)
            except Exception:
                dead.append(q)
        if dead:
            with self._lock:
                for q in dead:
                    if q in self._subscribers:
                        self._subscribers.remove(q)

    def iter_sse(self) -> Iterator[str]:
        q: queue.Queue[str | None] = queue.Queue(maxsize=32)
        with self._lock:
            self._subscribers.append(q)
        try:
            yield _sse_frame(self.event_dict())
            while True:
                try:
                    raw = q.get(timeout=_sse_heartbeat_sec())
                except queue.Empty:
                    yield ": heartbeat\n\n"
                    continue
                if raw is None:
                    break
                try:
                    data = json.loads(raw)
                except Exception:
                    data = self.event_dict()
                yield _sse_frame(data)
        finally:
            with self._lock:
                if q in self._subscribers:
                    self._subscribers.remove(q)


def _sse_frame(data: dict[str, Any]) -> str:
    body = json.dumps(data, ensure_ascii=False)
    return f"event: overview\ndata: {body}\n\n"


overview_store = OverviewStore()