liumangmang
157 lines
5.3 KiB
Python
157 lines
5.3 KiB
Python
"""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)
|