"""账户流水:后台定时拉取交易所 + SSE 版本推送.""" from __future__ import annotations import json import os import queue import threading import time from collections.abc import Callable, Iterator from datetime import datetime, timezone from typing import Any, Optional from lib.account_ledger.account_ledger_db import ( ensure_account_ledger_tables, meta_get, meta_set, prune_older_than, upsert_entries, ) ACCOUNT_LEDGER_POLL_SEC = float(os.getenv("ACCOUNT_LEDGER_POLL_SEC", "120")) ACCOUNT_LEDGER_LOOKBACK_DAYS = int(os.getenv("ACCOUNT_LEDGER_LOOKBACK_DAYS", "90")) ACCOUNT_LEDGER_SSE_HEARTBEAT_SEC = float(os.getenv("ACCOUNT_LEDGER_SSE_HEARTBEAT_SEC", "25")) class AccountLedgerStore: def __init__(self) -> None: self._lock = threading.Lock() self.version = 0 self._subscribers: list[queue.Queue[str | None]] = [] self._stop = threading.Event() self._thread: threading.Thread | None = None self._syncing = False self._get_db: Optional[Callable] = None self._fetch_fn: Optional[Callable[..., tuple[list[dict[str, Any]], list[str]]]] = None self._exchange_key = "" self.last_sync_at: Optional[float] = None self.last_error: str = "" self.last_upserted: int = 0 self._last_manual_at: float = 0.0 self._manual_cooldown_sec = float(os.getenv("ACCOUNT_LEDGER_MANUAL_COOLDOWN_SEC", "30")) def configure( self, *, get_db: Callable, fetch_fn: Callable[..., tuple[list[dict[str, Any]], list[str]]], exchange_key: str, ) -> None: self._get_db = get_db self._fetch_fn = fetch_fn self._exchange_key = (exchange_key or "").strip().lower() def start(self) -> None: if self._thread and self._thread.is_alive(): return if not self._get_db or not self._fetch_fn: return self._stop.clear() self._thread = threading.Thread( target=self._loop, daemon=True, name=f"account-ledger-{self._exchange_key or 'x'}" ) self._thread.start() def stop(self) -> None: self._stop.set() self._broadcast(close=True) def lookback_bounds_ms(self, start_ms: Optional[int] = None, end_ms: Optional[int] = None) -> tuple[int, int]: now = datetime.now(timezone.utc) end = int(end_ms) if end_ms is not None else int(now.timestamp() * 1000) floor = int(end - ACCOUNT_LEDGER_LOOKBACK_DAYS * 86400 * 1000) if start_ms is not None: start = max(int(start_ms), floor) else: start = floor if start > end: start, end = end, start return start, end def sync_once( self, *, reason: str = "poll", start_ms: Optional[int] = None, end_ms: Optional[int] = None, ) -> dict[str, Any]: if not self._get_db or not self._fetch_fn: return {"ok": False, "msg": "未配置"} with self._lock: if self._syncing: return {"ok": True, "busy": True, "ledger_version": self.version} if reason == "manual": gap = time.time() - self._last_manual_at if gap < self._manual_cooldown_sec: wait = int(self._manual_cooldown_sec - gap) + 1 return { "ok": False, "msg": f"同步过于频繁,请 {wait}s 后再试", "ledger_version": self.version, } self._syncing = True try: start, end = self.lookback_bounds_ms(start_ms, end_ms) rows, errors = self._fetch_fn(start_ms=start, end_ms=end) conn = self._get_db() try: ensure_account_ledger_tables(conn) n = upsert_entries(conn, rows or []) # 保留略宽于 lookback 的缓存 prune_ms = int( (datetime.now(timezone.utc).timestamp() - (ACCOUNT_LEDGER_LOOKBACK_DAYS + 7) * 86400) * 1000 ) prune_older_than(conn, prune_ms) self.last_sync_at = time.time() self.last_upserted = n self.last_error = "; ".join(errors[:3]) if errors else "" meta_set(conn, "last_sync_at", str(self.last_sync_at)) meta_set(conn, "last_error", self.last_error) meta_set(conn, "last_upserted", str(n)) conn.commit() finally: try: conn.close() except Exception: pass if reason == "manual": self._last_manual_at = time.time() ver = self.bump(reason) return { "ok": True, "ledger_version": ver, "upserted": n, "errors": errors, "start_ms": start, "end_ms": end, } except Exception as e: self.last_error = str(e) try: conn = self._get_db() try: ensure_account_ledger_tables(conn) meta_set(conn, "last_error", self.last_error) conn.commit() finally: conn.close() except Exception: pass return {"ok": False, "msg": str(e), "ledger_version": self.version} finally: with self._lock: self._syncing = False def bump(self, reason: str = "poll") -> int: with self._lock: self.version += 1 ver = self.version payload = json.dumps( {"ledger_version": ver, "reason": reason, "exchange": self._exchange_key}, ensure_ascii=False, ) self._broadcast(payload) return ver def status_dict(self) -> dict[str, Any]: last_at = self.last_sync_at if last_at is None and self._get_db: try: conn = self._get_db() try: ensure_account_ledger_tables(conn) raw = meta_get(conn, "last_sync_at", "") if raw: last_at = float(raw) self.last_error = meta_get(conn, "last_error", self.last_error) finally: conn.close() except Exception: pass return { "ledger_version": self.version, "poll_sec": ACCOUNT_LEDGER_POLL_SEC, "lookback_days": ACCOUNT_LEDGER_LOOKBACK_DAYS, "last_sync_at": last_at, "last_error": self.last_error, "last_upserted": self.last_upserted, "exchange": self._exchange_key, } def _loop(self) -> None: # 启动后稍等再拉,避免和启动高峰撞车 if self._stop.wait(3): return while not self._stop.is_set(): try: self.sync_once(reason="poll") except Exception: pass if self._stop.wait(ACCOUNT_LEDGER_POLL_SEC): break def _broadcast(self, event: str | None = None, *, close: bool = False) -> None: with self._lock: subs = list(self._subscribers) dead: list[queue.Queue[str | None]] = [] for q in subs: try: q.put_nowait(None if close else event) except Exception: dead.append(q) if dead: with self._lock: for q in dead: if q in self._subscribers: self._subscribers.remove(q) def _subscribe(self) -> queue.Queue[str | None]: q: queue.Queue[str | None] = queue.Queue(maxsize=16) with self._lock: self._subscribers.append(q) return q def _unsubscribe(self, q: queue.Queue[str | None]) -> None: with self._lock: if q in self._subscribers: self._subscribers.remove(q) def iter_sse(self) -> Iterator[str]: q = self._subscribe() try: yield f"event: ledger\ndata: {json.dumps({'ledger_version': self.version, 'reason': 'hello'}, ensure_ascii=False)}\n\n" last_hb = time.time() while not self._stop.is_set(): try: item = q.get(timeout=1.0) except queue.Empty: item = "timeout" if item is None: break if item != "timeout": yield f"event: ledger\ndata: {item}\n\n" last_hb = time.time() elif time.time() - last_hb >= ACCOUNT_LEDGER_SSE_HEARTBEAT_SEC: yield ": heartbeat\n\n" last_hb = time.time() finally: self._unsubscribe(q) account_ledger_store = AccountLedgerStore()