Initial commit

2026-05-12 17:51:53 +08:00
commit b564ca4797
55 changed files with 6407 additions and 0 deletions
@@ -0,0 +1,156 @@
+"""APScheduler background scheduler for upstream checks."""
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime, timezone
+
+from apscheduler.schedulers.background import BackgroundScheduler
+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.upstream_client import UpstreamClient, UpstreamError, build_snapshot
+from app.services.snapshot_service import diff_snapshots
+from app.services import webhook_service
+from app.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+_scheduler = BackgroundScheduler(timezone="UTC")
+
+
+def get_scheduler() -> BackgroundScheduler:
+    return _scheduler
+
+
+def _check_upstream(upstream_id: int) -> None:
+    """Full upstream check executed by scheduler (runs in thread)."""
+    settings = get_settings()
+    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 "{}")
+        client = 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),
+        )
+
+        was_unhealthy = upstream.last_status == "unhealthy"
+
+        try:
+            client.login()
+            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
+            )
+        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
+            if upstream.consecutive_failures >= threshold and upstream.last_status != "unhealthy":
+                upstream.last_status = "unhealthy"
+                db.commit()
+                webhook_service.send_status_event(
+                    db, upstream.id, upstream.name, upstream.base_url,
+                    "upstream_unhealthy", str(exc)
+                )
+            else:
+                db.commit()
+            logger.warning("upstream %s check failed: %s", upstream.name, exc)
+            return
+
+        # success path
+        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)
+
+        # 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()
+
+        if was_unhealthy:
+            webhook_service.send_status_event(
+                db, upstream.id, upstream.name, upstream.base_url, "upstream_recovered"
+            )
+
+        if changes:
+            webhook_service.send_rate_changed(
+                db, upstream.id, upstream.name, upstream.base_url, changes
+            )
+            logger.info("upstream %s: %d rate change(s)", upstream.name, len(changes))
+        else:
+            logger.debug("upstream %s: no changes", upstream.name)
+
+    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,
+    )
+    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)
+        logger.info("scheduler started with %d upstream job(s)", len(upstreams))
+    finally:
+        db.close()
+
+
+def stop_scheduler() -> None:
+    if _scheduler.running:
+        _scheduler.shutdown(wait=False)