192 lines
7.0 KiB
Python
192 lines
7.0 KiB
Python
"""监控调度器：线程池并发轮询所有启用账号。
|
|
|
|
调度逻辑：
|
|
- 每个账号记录 next_check_at（epoch 秒），间隔取平台覆盖值或全局默认
|
|
- 调度线程每秒扫描一次，到期的账号且无进行中任务则提交线程池
|
|
- 任务完成后把 next_check_at 更新为 now + interval，避免重复执行
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
|
|
from app import db
|
|
from app.alert import evaluate, notify_disabled
|
|
from app.config import Config
|
|
from app.fetcher import fetch_balance
|
|
|
|
logger = logging.getLogger("monitor.scheduler")
|
|
|
|
# JOIN 查询后拆分为账号/平台两个 dict 用到的字段
|
|
_ACCOUNT_FIELDS = [
|
|
"id", "platform_id", "name", "api_key", "threshold", "enabled",
|
|
"alert_armed", "last_balance", "last_status", "last_error",
|
|
"last_check_at", "note",
|
|
]
|
|
_PLATFORM_FIELDS = [
|
|
"platform_id", "platform_name", "currency", "icon", "method", "url",
|
|
"headers", "body", "balance_path", "interval_seconds",
|
|
"retry_count", "timeout_seconds", "platform_enabled", "note",
|
|
]
|
|
|
|
_ACCOUNT_SQL = """
|
|
SELECT a.*,
|
|
p.name AS platform_name, p.currency, p.icon, p.method, p.url,
|
|
p.headers, p.body, p.balance_path, p.interval_seconds,
|
|
p.retry_count, p.timeout_seconds, p.enabled AS platform_enabled,
|
|
p.note AS platform_note
|
|
FROM accounts a JOIN platforms p ON p.id = a.platform_id
|
|
"""
|
|
|
|
|
|
def _split(row: dict) -> tuple[dict, dict]:
|
|
"""把 JOIN 行拆成 (account, platform) 两个 dict。"""
|
|
account = {k: row[k] for k in _ACCOUNT_FIELDS if k in row}
|
|
platform = {k: row[k] for k in _PLATFORM_FIELDS if k in row}
|
|
platform["name"] = platform.pop("platform_name")
|
|
platform["id"] = platform.pop("platform_id")
|
|
platform["enabled"] = platform.pop("platform_enabled")
|
|
return account, platform
|
|
|
|
|
|
def _account_interval(platform: dict, cfg: Config) -> float:
|
|
return float(platform.get("interval_seconds") or cfg.global_interval_seconds)
|
|
|
|
|
|
class Monitor:
|
|
def __init__(self, cfg: Config) -> None:
|
|
self.cfg = cfg
|
|
self._stop = threading.Event()
|
|
self._pool: ThreadPoolExecutor | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self._next_check: dict[int, float] = {}
|
|
self._inflight: set[int] = set()
|
|
self._lock = threading.Lock()
|
|
|
|
# ---------- 生命周期 ----------
|
|
|
|
def start(self) -> None:
|
|
if self._thread is not None:
|
|
return
|
|
self._pool = ThreadPoolExecutor(max_workers=self.cfg.max_workers, thread_name_prefix="check")
|
|
self._thread = threading.Thread(target=self._loop, name="monitor-loop", daemon=True)
|
|
self._thread.start()
|
|
logger.info(
|
|
"监控调度已启动，全局间隔 %ss，并发 %s",
|
|
self.cfg.global_interval_seconds,
|
|
self.cfg.max_workers,
|
|
)
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=5)
|
|
if self._pool is not None:
|
|
self._pool.shutdown(wait=False, cancel_futures=True)
|
|
logger.info("监控调度已停止")
|
|
|
|
# ---------- 调度循环 ----------
|
|
|
|
def _loop(self) -> None:
|
|
while not self._stop.is_set():
|
|
try:
|
|
self._dispatch()
|
|
except Exception:
|
|
logger.exception("调度循环异常")
|
|
self._stop.wait(1)
|
|
|
|
def _dispatch(self) -> None:
|
|
with db.get_conn() as conn:
|
|
rows = conn.execute(
|
|
_ACCOUNT_SQL + " WHERE a.enabled = 1 AND p.enabled = 1"
|
|
).fetchall()
|
|
now = time.time()
|
|
for row in rows:
|
|
data = dict(row)
|
|
aid = data["id"]
|
|
with self._lock:
|
|
if aid in self._inflight:
|
|
continue
|
|
if now < self._next_check.get(aid, 0):
|
|
continue
|
|
self._inflight.add(aid)
|
|
assert self._pool is not None
|
|
future = self._pool.submit(self._check_one, data)
|
|
future.add_done_callback(lambda f, a=aid: self._on_done(a, f))
|
|
|
|
def _on_done(self, account_id: int, future: Future) -> None:
|
|
with self._lock:
|
|
self._inflight.discard(account_id)
|
|
if future.exception():
|
|
logger.error("账号 %s 检查任务异常: %s", account_id, future.exception())
|
|
|
|
# ---------- 单次检查 ----------
|
|
|
|
def _check_one(self, row: dict) -> None:
|
|
cfg = self.cfg
|
|
account, platform = _split(row)
|
|
retry = platform.get("retry_count")
|
|
if retry is None:
|
|
retry = cfg.retry_count
|
|
timeout = platform.get("timeout_seconds")
|
|
if timeout is None:
|
|
timeout = cfg.timeout_seconds
|
|
result = fetch_balance(platform, account["api_key"], int(retry), int(timeout))
|
|
|
|
with db.get_conn() as conn:
|
|
if result.ok:
|
|
conn.execute(
|
|
"UPDATE accounts SET last_balance=?, last_status='ok', last_error='', "
|
|
"last_check_at=datetime('now','localtime') WHERE id=?",
|
|
(result.balance, account["id"]),
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO balance_history (account_id, balance) VALUES (?, ?)",
|
|
(account["id"], result.balance),
|
|
)
|
|
elif result.auth_error:
|
|
conn.execute(
|
|
"UPDATE accounts SET enabled=0, last_status='disabled', last_error=?, "
|
|
"last_check_at=datetime('now','localtime') WHERE id=?",
|
|
(result.error, account["id"]),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"UPDATE accounts SET last_status='error', last_error=?, "
|
|
"last_check_at=datetime('now','localtime') WHERE id=?",
|
|
(result.error, account["id"]),
|
|
)
|
|
|
|
if result.auth_error:
|
|
notify_disabled(account, platform, cfg)
|
|
elif result.ok:
|
|
evaluate(account, platform, result.balance, cfg)
|
|
|
|
# 调度：无论成功失败，按间隔排下一次
|
|
with self._lock:
|
|
self._next_check[account["id"]] = time.time() + _account_interval(platform, cfg)
|
|
|
|
# ---------- 手动触发 ----------
|
|
|
|
def check_now(self, account_id: int) -> bool:
|
|
"""API 手动立即检查；账号/平台未启用或已在检查中则返回 False。"""
|
|
with db.get_conn() as conn:
|
|
row = conn.execute(
|
|
_ACCOUNT_SQL + " WHERE a.id = ? AND a.enabled = 1 AND p.enabled = 1",
|
|
(account_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return False
|
|
data = dict(row)
|
|
with self._lock:
|
|
if account_id in self._inflight:
|
|
return False
|
|
self._inflight.add(account_id)
|
|
assert self._pool is not None
|
|
future = self._pool.submit(self._check_one, data)
|
|
future.add_done_callback(lambda f, a=account_id: self._on_done(a, f))
|
|
return True
|