112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from app.extensions import db
|
|
from app.models import ApplyLog, Mirror, Profile, ProfileItem, Setting, SystemState
|
|
from app.models import utcnow
|
|
from app.services.applier import apply_profile, rollback_apply
|
|
from app.services.tester import get_best_mirror, test_mirror
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_scheduler = None
|
|
|
|
|
|
def _auto_switch_enabled(category: str) -> bool:
|
|
return Setting.get(f"auto_switch_{category}", "false").lower() == "true"
|
|
|
|
|
|
def _get_switch_count_today() -> int:
|
|
key = f"switch_count_{utcnow().strftime('%Y%m%d')}"
|
|
return int(Setting.get(key, "0") or "0")
|
|
|
|
|
|
def _increment_switch_count() -> None:
|
|
key = f"switch_count_{utcnow().strftime('%Y%m%d')}"
|
|
Setting.set(key, str(_get_switch_count_today() + 1))
|
|
|
|
|
|
def check_pending_rollbacks(app) -> None:
|
|
with app.app_context():
|
|
now = utcnow()
|
|
pending = ApplyLog.query.filter(
|
|
ApplyLog.status == "pending_confirm",
|
|
ApplyLog.confirmed.is_(False),
|
|
ApplyLog.rollback_at.isnot(None),
|
|
ApplyLog.rollback_at <= now,
|
|
).all()
|
|
|
|
for log in pending:
|
|
logger.info("Auto rollback for apply log %s", log.id)
|
|
rollback_apply(log.id)
|
|
|
|
|
|
def run_auto_switch(app) -> None:
|
|
with app.app_context():
|
|
max_switches = int(Setting.get("max_switches_per_day", "3") or "3")
|
|
if _get_switch_count_today() >= max_switches:
|
|
logger.info("Max switches per day reached")
|
|
return
|
|
|
|
state = SystemState.query.first()
|
|
if not state or not state.current_profile_id:
|
|
return
|
|
|
|
profile = Profile.query.get(state.current_profile_id)
|
|
if not profile:
|
|
return
|
|
|
|
for category in ("dns", "apt", "docker", "github", "pip", "npm"):
|
|
if not _auto_switch_enabled(category):
|
|
continue
|
|
|
|
item = ProfileItem.query.filter_by(profile_id=profile.id, category=category).first()
|
|
if not item or not item.mirror:
|
|
continue
|
|
|
|
results = test_mirror(item.mirror)
|
|
if results and all(r.success for r in results):
|
|
continue
|
|
|
|
best = get_best_mirror(category)
|
|
if not best or best.id == item.mirror_id:
|
|
if Setting.get("rollback_on_all_fail", "true").lower() == "true":
|
|
logger.warning("All mirrors failed for %s", category)
|
|
continue
|
|
|
|
item.mirror_id = best.id
|
|
db.session.commit()
|
|
apply_profile(profile)
|
|
_increment_switch_count()
|
|
logger.info("Auto-switched %s to %s", category, best.name)
|
|
break
|
|
|
|
|
|
def init_scheduler(app) -> None:
|
|
global _scheduler
|
|
from apscheduler.schedulers.background import BackgroundScheduler
|
|
|
|
if _scheduler is not None:
|
|
return
|
|
|
|
interval = int(app.config.get("AUTO_TEST_INTERVAL_MINUTES", 30))
|
|
_scheduler = BackgroundScheduler(daemon=True)
|
|
_scheduler.add_job(
|
|
check_pending_rollbacks,
|
|
"interval",
|
|
minutes=1,
|
|
args=[app],
|
|
id="rollback_checker",
|
|
)
|
|
_scheduler.add_job(
|
|
run_auto_switch,
|
|
"interval",
|
|
minutes=interval,
|
|
args=[app],
|
|
id="auto_switch",
|
|
)
|
|
_scheduler.start()
|
|
logger.info("Scheduler started (interval=%s min)", interval)
|