dekun
253 lines
8.9 KiB
Python
253 lines
8.9 KiB
Python
"""账户流水:后台定时拉取交易所 + 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()
|