479 lines
18 KiB
Python
479 lines
18 KiB
Python
"""APScheduler background scheduler for upstream checks and daily tasks."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
from apscheduler.executors.pool import ThreadPoolExecutor
|
|
from apscheduler.schedulers.background import BackgroundScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.database import SessionLocal
|
|
from app.models.upstream import Upstream
|
|
from app.models.snapshot import UpstreamRateSnapshot
|
|
from app.services.auth_config import normalize_auth_config
|
|
from app.services.upstream_client import UpstreamClient, build_snapshot
|
|
from app.services.snapshot_service import diff_snapshots, prune_snapshots, write_balance_snapshot
|
|
from app.services import webhook_service
|
|
from app.services import website_sync
|
|
from app.config import get_settings
|
|
from app.routers.external_api_logs import clean_expired_logs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_scheduler = BackgroundScheduler(timezone="UTC", executors={"default": ThreadPoolExecutor(max_workers=1)})
|
|
|
|
|
|
def get_scheduler() -> BackgroundScheduler:
|
|
return _scheduler
|
|
|
|
|
|
def _persist_auth_config_update(upstream: Upstream, updated_config: dict[str, Any]) -> None:
|
|
config_json = json.dumps(
|
|
normalize_auth_config(upstream.auth_type, updated_config),
|
|
ensure_ascii=False,
|
|
)
|
|
updated_at = datetime.now(timezone.utc)
|
|
db = SessionLocal()
|
|
try:
|
|
row = db.query(Upstream).filter(Upstream.id == upstream.id).first()
|
|
if row:
|
|
row.auth_config_json = config_json
|
|
row.updated_at = updated_at
|
|
db.commit()
|
|
upstream.auth_config_json = row.auth_config_json
|
|
upstream.updated_at = updated_at
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _check_upstream(upstream_id: int) -> None:
|
|
"""Full upstream check executed by scheduler (runs in thread).
|
|
|
|
Phase 1 — upstream API call + snapshot write (single transaction).
|
|
Phase 2 — webhook/website sync (separate sessions, so a notification
|
|
failure never rolls back the snapshot).
|
|
"""
|
|
settings = get_settings()
|
|
# ── Phase 1: upstream check + DB write ──────────────────────────
|
|
db: Session = SessionLocal()
|
|
try:
|
|
upstream = db.query(Upstream).filter(Upstream.id == upstream_id).first()
|
|
if not upstream or not upstream.enabled:
|
|
_remove_job(upstream_id)
|
|
return
|
|
|
|
auth_config = json.loads(upstream.auth_config_json or "{}")
|
|
was_unhealthy = upstream.last_status == "unhealthy"
|
|
balance_alert_triggered = False
|
|
snapshot = None
|
|
changes = None
|
|
|
|
with UpstreamClient(
|
|
base_url=upstream.base_url,
|
|
api_prefix=upstream.api_prefix,
|
|
auth_type=upstream.auth_type,
|
|
auth_config=auth_config,
|
|
timeout=float(upstream.timeout_seconds),
|
|
on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated),
|
|
target_id=upstream.id,
|
|
target_name=upstream.name,
|
|
) as client:
|
|
try:
|
|
client.ensure_authenticated()
|
|
groups = client.get_available_groups(upstream.groups_endpoint)
|
|
raw_rates = client.get_group_rates(upstream.rate_endpoint)
|
|
snapshot = build_snapshot(
|
|
upstream.id, upstream.base_url, upstream.api_prefix, groups, raw_rates
|
|
)
|
|
# ── Balance fetch (inside with block, client still open) ──
|
|
balance: Optional[float] = None
|
|
if upstream.balance_endpoint and upstream.balance_response_path:
|
|
try:
|
|
raw_balance = client.get_balance(upstream.balance_endpoint, upstream.balance_response_path)
|
|
if raw_balance is not None:
|
|
divisor = upstream.balance_divisor or 1.0
|
|
balance = raw_balance / divisor
|
|
except Exception as exc:
|
|
logger.warning("upstream %s balance fetch failed: %s", upstream.name, exc)
|
|
if balance is not None:
|
|
upstream.balance = balance
|
|
upstream.balance_updated_at = datetime.now(timezone.utc)
|
|
write_balance_snapshot(db, upstream.id, balance)
|
|
# ── 余额告警阈值检查 ──
|
|
threshold = upstream.balance_alert_threshold
|
|
if threshold is not None and threshold > 0:
|
|
if balance < threshold and not upstream.balance_alert_notified:
|
|
upstream.balance_alert_notified = True
|
|
balance_alert_triggered = True
|
|
elif balance >= threshold and upstream.balance_alert_notified:
|
|
upstream.balance_alert_notified = False
|
|
except Exception as exc:
|
|
# failure path
|
|
upstream.consecutive_failures = (upstream.consecutive_failures or 0) + 1
|
|
upstream.last_error = str(exc)
|
|
upstream.last_checked_at = datetime.now(timezone.utc)
|
|
threshold = settings.unhealthy_threshold
|
|
became_unhealthy = (
|
|
upstream.consecutive_failures >= threshold
|
|
and upstream.last_status != "unhealthy"
|
|
)
|
|
if became_unhealthy:
|
|
upstream.last_status = "unhealthy"
|
|
db.commit()
|
|
logger.warning("upstream %s check failed: %s", upstream.name, exc)
|
|
# Phase 2: notify unhealthy in a fresh session
|
|
if became_unhealthy:
|
|
_notify_status(upstream.id, upstream.name, upstream.base_url,
|
|
"upstream_unhealthy", str(exc))
|
|
return
|
|
|
|
# success path (client auto-closed by `with`)
|
|
|
|
prev_snapshot_row = (
|
|
db.query(UpstreamRateSnapshot)
|
|
.filter(UpstreamRateSnapshot.upstream_id == upstream_id)
|
|
.order_by(UpstreamRateSnapshot.captured_at.desc())
|
|
.first()
|
|
)
|
|
previous = json.loads(prev_snapshot_row.snapshot_json) if prev_snapshot_row else None
|
|
changes = diff_snapshots(previous, snapshot)
|
|
|
|
# save new snapshot
|
|
new_row = UpstreamRateSnapshot(
|
|
upstream_id=upstream_id,
|
|
snapshot_json=json.dumps(snapshot, ensure_ascii=False),
|
|
captured_at=datetime.now(timezone.utc),
|
|
)
|
|
db.add(new_row)
|
|
prune_snapshots(db, upstream_id, settings.snapshot_retention_count)
|
|
|
|
# update upstream status
|
|
upstream.last_status = "healthy"
|
|
upstream.last_checked_at = datetime.now(timezone.utc)
|
|
upstream.last_error = None
|
|
upstream.consecutive_failures = 0
|
|
db.commit()
|
|
|
|
logger.info(
|
|
"upstream %s: %d rate change(s)" if changes else "upstream %s: no changes",
|
|
upstream.name, len(changes) if changes else 0,
|
|
)
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
# ── Phase 2: key sync (independent session) ───────────────────
|
|
if snapshot:
|
|
captured_at = snapshot.get("captured_at")
|
|
if isinstance(captured_at, str):
|
|
from datetime import datetime as dt
|
|
try:
|
|
captured_at = dt.fromisoformat(captured_at)
|
|
except Exception:
|
|
captured_at = datetime.now(timezone.utc)
|
|
elif captured_at is None:
|
|
captured_at = datetime.now(timezone.utc)
|
|
_sync_upstream_keys(upstream_id, snapshot, captured_at)
|
|
|
|
# ── Phase 3: notifications (independent sessions) ──────────────
|
|
if was_unhealthy:
|
|
_notify_status(upstream_id, upstream.name, upstream.base_url, "upstream_recovered")
|
|
|
|
if changes:
|
|
_notify_rate_changed(upstream_id, upstream.name, upstream.base_url, changes)
|
|
_sync_website_bindings(upstream_id, changes)
|
|
_sync_account_priorities(upstream_id)
|
|
|
|
if balance_alert_triggered:
|
|
_notify_balance_low(
|
|
upstream_id, upstream.name, upstream.base_url,
|
|
upstream.balance, upstream.balance_alert_threshold,
|
|
)
|
|
|
|
|
|
def _notify_status(
|
|
upstream_id: int,
|
|
upstream_name: str,
|
|
base_url: str,
|
|
event: str,
|
|
error: str = "",
|
|
) -> None:
|
|
db = SessionLocal()
|
|
try:
|
|
webhook_service.send_status_event(db, upstream_id, upstream_name, base_url, event, error)
|
|
except Exception:
|
|
logger.exception("status webhook failed for upstream %s", upstream_name)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _notify_rate_changed(
|
|
upstream_id: int,
|
|
upstream_name: str,
|
|
base_url: str,
|
|
changes: list[dict[str, Any]],
|
|
) -> None:
|
|
db = SessionLocal()
|
|
try:
|
|
webhook_service.send_rate_changed(db, upstream_id, upstream_name, base_url, changes)
|
|
except Exception:
|
|
logger.exception("rate webhook failed for upstream %s", upstream_name)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _notify_balance_low(
|
|
upstream_id: int,
|
|
upstream_name: str,
|
|
base_url: str,
|
|
balance: float,
|
|
threshold: float,
|
|
) -> None:
|
|
db = SessionLocal()
|
|
try:
|
|
webhook_service.send_balance_low(db, upstream_id, upstream_name, base_url, balance, threshold)
|
|
except Exception:
|
|
logger.exception("balance low webhook failed for upstream %s", upstream_name)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _sync_upstream_keys(upstream_id: int, snapshot: dict[str, Any], captured_at: datetime) -> None:
|
|
"""上游检测成功后同步 SmartUp Key 状态（远端删除/分组删除）。
|
|
|
|
委托给 website_sync.reconcile_upstream_keys 实现核心逻辑。
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
active_group_ids = set(snapshot.get("groups", {}).keys())
|
|
|
|
# 获取倒数第二个快照作为 previous snapshot
|
|
prev_active_group_ids: set[str] | None = None
|
|
from app.models.snapshot import UpstreamRateSnapshot
|
|
snapshots = (
|
|
db.query(UpstreamRateSnapshot)
|
|
.filter(UpstreamRateSnapshot.upstream_id == upstream_id)
|
|
.order_by(UpstreamRateSnapshot.captured_at.desc())
|
|
.limit(2)
|
|
.all()
|
|
)
|
|
if len(snapshots) >= 2:
|
|
prev_snapshot = snapshots[1]
|
|
try:
|
|
prev_data = json.loads(prev_snapshot.snapshot_json or "{}")
|
|
prev_groups = prev_data.get("groups") or {}
|
|
if isinstance(prev_groups, dict):
|
|
prev_active_group_ids = set(prev_groups.keys())
|
|
except Exception:
|
|
pass
|
|
|
|
# 用 UpstreamClient 查询远端活跃 Key ID 集合
|
|
remote_key_ids: set[str] | None = None
|
|
try:
|
|
upstream = db.query(Upstream).filter(Upstream.id == upstream_id).first()
|
|
if upstream:
|
|
auth_config = json.loads(upstream.auth_config_json or "{}")
|
|
with UpstreamClient(
|
|
base_url=upstream.base_url,
|
|
api_prefix=upstream.api_prefix,
|
|
auth_type=upstream.auth_type,
|
|
auth_config=auth_config,
|
|
timeout=float(upstream.timeout_seconds),
|
|
on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated),
|
|
target_id=upstream.id,
|
|
target_name=upstream.name,
|
|
) as client:
|
|
client.ensure_authenticated()
|
|
remote_key_ids = website_sync._fetch_remote_managed_key_ids(db, client, upstream_id)
|
|
except Exception as exc:
|
|
logger.warning("sync upstream keys list failed for %s: %s", upstream_id, exc)
|
|
|
|
website_sync.reconcile_upstream_keys(db, upstream_id, active_group_ids, remote_key_ids, captured_at, prev_active_group_ids=prev_active_group_ids)
|
|
db.commit()
|
|
except Exception:
|
|
logger.exception("key sync failed for upstream %s", upstream_id)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _sync_account_priorities(upstream_id: int) -> None:
|
|
"""倍率变更后自动更新已导入下游账号的 priority。"""
|
|
db = SessionLocal()
|
|
try:
|
|
website_sync.sync_account_priorities_for_upstream(db, upstream_id)
|
|
except Exception:
|
|
logger.exception("account priority sync failed for upstream %s", upstream_id)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _sync_website_bindings(upstream_id: int, changes: list[dict[str, Any]]) -> None:
|
|
db = SessionLocal()
|
|
try:
|
|
website_sync.sync_affected_bindings(db, upstream_id, changes)
|
|
except Exception:
|
|
logger.exception("website sync failed for upstream %s", upstream_id)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _remove_job(upstream_id: int) -> None:
|
|
job_id = f"upstream_{upstream_id}"
|
|
if _scheduler.get_job(job_id):
|
|
_scheduler.remove_job(job_id)
|
|
|
|
|
|
def refresh_upstream(upstream_id: int, interval_seconds: int = 0, enabled: bool = True) -> None:
|
|
"""Add/update/remove a scheduler job for the given upstream."""
|
|
job_id = f"upstream_{upstream_id}"
|
|
if not enabled or interval_seconds <= 0:
|
|
_remove_job(upstream_id)
|
|
return
|
|
_scheduler.add_job(
|
|
_check_upstream,
|
|
"interval",
|
|
seconds=interval_seconds,
|
|
id=job_id,
|
|
args=[upstream_id],
|
|
replace_existing=True,
|
|
coalesce=True,
|
|
max_instances=1,
|
|
misfire_grace_time=60,
|
|
jitter=30,
|
|
)
|
|
logger.info("scheduler job %s set to %ds interval", job_id, interval_seconds)
|
|
|
|
|
|
def start_scheduler() -> None:
|
|
"""Start scheduler and load all enabled upstreams."""
|
|
_scheduler.start()
|
|
db: Session = SessionLocal()
|
|
try:
|
|
upstreams = db.query(Upstream).filter(Upstream.enabled == True).all()
|
|
for u in upstreams:
|
|
refresh_upstream(u.id, u.check_interval_seconds, u.enabled)
|
|
# Daily cleanup of expired external API logs
|
|
_scheduler.add_job(
|
|
_cleanup_external_logs,
|
|
"interval",
|
|
hours=24,
|
|
id="cleanup_external_api_logs",
|
|
replace_existing=True,
|
|
coalesce=True,
|
|
max_instances=1,
|
|
misfire_grace_time=3600,
|
|
)
|
|
# Daily finance summary at 00:05 Asia/Shanghai
|
|
_scheduler.add_job(
|
|
_compute_finance_daily_summary,
|
|
trigger=CronTrigger(timezone="Asia/Shanghai", hour=0, minute=5),
|
|
id="finance_daily_summary",
|
|
replace_existing=True,
|
|
coalesce=True,
|
|
max_instances=1,
|
|
misfire_grace_time=3600,
|
|
)
|
|
# Daily balance snapshot capture at 23:59 Asia/Shanghai
|
|
_scheduler.add_job(
|
|
_capture_balance_snapshots,
|
|
trigger=CronTrigger(timezone="Asia/Shanghai", hour=23, minute=59),
|
|
id="balance_snapshot_daily",
|
|
replace_existing=True,
|
|
coalesce=True,
|
|
max_instances=1,
|
|
misfire_grace_time=600,
|
|
)
|
|
logger.info("scheduler started with %d upstream job(s)", len(upstreams))
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _cleanup_external_logs() -> None:
|
|
"""Delete external API logs older than the configured retention period."""
|
|
db = SessionLocal()
|
|
try:
|
|
deleted = clean_expired_logs(db)
|
|
if deleted:
|
|
logger.info("cleaned %d expired external API log(s)", deleted)
|
|
except Exception:
|
|
logger.exception("failed to clean expired external API logs")
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _compute_finance_daily_summary() -> None:
|
|
"""Compute and save finance daily summary for yesterday (Asia/Shanghai).
|
|
|
|
Runs daily at 00:05 Asia/Shanghai. Delegates to save_summary_if_absent
|
|
which handles the skip-if-exists logic and concurrent-save race.
|
|
"""
|
|
from app.services.finance_service import yesterday_shanghai, save_summary_if_absent
|
|
|
|
target = yesterday_shanghai()
|
|
db = SessionLocal()
|
|
try:
|
|
saved = save_summary_if_absent(db, target)
|
|
if saved:
|
|
logger.info("finance daily summary for %s saved", target)
|
|
else:
|
|
logger.info("finance daily summary for %s already exists, skipping", target)
|
|
except Exception:
|
|
logger.exception("failed to compute finance daily summary for %s", target)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _capture_balance_snapshots() -> None:
|
|
"""Fetch balance for all balance_delta-mode upstreams and write a snapshot.
|
|
|
|
Runs daily at 23:59 Asia/Shanghai to capture the end-of-day boundary balance.
|
|
Only writes a snapshot if balance fetch succeeds. Does NOT update the
|
|
upstream.balance field (that's the regular check's job).
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
upstreams = db.query(Upstream).filter(
|
|
Upstream.enabled == True,
|
|
Upstream.finance_cost_mode == "balance_delta",
|
|
Upstream.balance_endpoint != "",
|
|
Upstream.balance_response_path != "",
|
|
).all()
|
|
for u in upstreams:
|
|
try:
|
|
auth_config = json.loads(u.auth_config_json or "{}")
|
|
with UpstreamClient(
|
|
base_url=u.base_url,
|
|
api_prefix=u.api_prefix,
|
|
auth_type=u.auth_type,
|
|
auth_config=auth_config,
|
|
timeout=float(u.timeout_seconds),
|
|
on_auth_config_update=lambda updated: _persist_auth_config_update(u, updated),
|
|
target_id=u.id,
|
|
target_name=u.name,
|
|
) as client:
|
|
client.ensure_authenticated()
|
|
raw_balance = client.get_balance(u.balance_endpoint, u.balance_response_path)
|
|
if raw_balance is not None:
|
|
balance = raw_balance / (u.balance_divisor or 1.0)
|
|
write_balance_snapshot(db, u.id, balance)
|
|
db.commit()
|
|
logger.info("balance boundary snapshot for upstream %s: %.4f", u.name, balance)
|
|
else:
|
|
logger.warning("balance boundary snapshot: upstream %s returned None", u.name)
|
|
except Exception as exc:
|
|
db.rollback()
|
|
logger.warning("balance boundary snapshot failed for upstream %s: %s", u.name, exc)
|
|
except Exception:
|
|
logger.exception("failed to capture balance boundary snapshots")
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def stop_scheduler() -> None:
|
|
if _scheduler.running:
|
|
_scheduler.shutdown(wait=True)
|