95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
"""轻量级调度中心。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker
|
|
|
|
from ..core.logging import get_logger
|
|
from . import workload_service
|
|
from . import workload_alert_service
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class WorkloadRecalcScheduler:
|
|
"""每日 00:00 重算 workload_ratio。"""
|
|
|
|
def __init__(self, session_factory: async_sessionmaker):
|
|
self._session_factory = session_factory
|
|
self._task: Optional[asyncio.Task] = None
|
|
self._stopped = asyncio.Event()
|
|
|
|
def start(self) -> None:
|
|
if self._task and not self._task.done():
|
|
return
|
|
self._stopped.clear()
|
|
self._task = asyncio.create_task(self._run_loop(), name="workload-recalc-daily")
|
|
logger.info("workload_scheduler_started")
|
|
|
|
async def stop(self) -> None:
|
|
self._stopped.set()
|
|
if not self._task:
|
|
return
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
logger.info("workload_scheduler_stopped")
|
|
|
|
async def run_once(self) -> dict:
|
|
"""手动执行一次，便于调试与测试。"""
|
|
return await self._run_recalc()
|
|
|
|
async def _run_loop(self) -> None:
|
|
while not self._stopped.is_set():
|
|
seconds = self._seconds_until_next_midnight()
|
|
try:
|
|
await asyncio.wait_for(self._stopped.wait(), timeout=seconds)
|
|
break
|
|
except TimeoutError:
|
|
pass
|
|
|
|
try:
|
|
await self._run_recalc()
|
|
except Exception as exc: # pragma: no cover
|
|
logger.warning("workload_scheduler_recalc_failed", error=str(exc))
|
|
|
|
@staticmethod
|
|
def _seconds_until_next_midnight() -> float:
|
|
now = datetime.now()
|
|
tomorrow = (now + timedelta(days=1)).date()
|
|
next_midnight = datetime.combine(tomorrow, datetime.min.time())
|
|
return max((next_midnight - now).total_seconds(), 1.0)
|
|
|
|
async def _run_recalc(self) -> dict:
|
|
async with self._session_factory() as db:
|
|
tenant_ids = await workload_service.list_tenant_ids(db)
|
|
total_users = 0
|
|
total_suggestions = 0
|
|
for tenant_id in tenant_ids or [1]:
|
|
total_users += await workload_service.recalculate_all_users_workload_ratio(
|
|
db,
|
|
tenant_id=tenant_id,
|
|
)
|
|
summary = await workload_alert_service.process_tenant_unresolved_overloads(
|
|
db,
|
|
tenant_id=tenant_id,
|
|
)
|
|
total_suggestions += int(summary.get("suggestion_count", 0))
|
|
logger.info(
|
|
"workload_scheduler_recalc_done",
|
|
tenant_count=len(tenant_ids or [1]),
|
|
user_count=total_users,
|
|
suggestion_count=total_suggestions,
|
|
)
|
|
return {
|
|
"tenant_count": len(tenant_ids or [1]),
|
|
"user_count": total_users,
|
|
"suggestion_count": total_suggestions,
|
|
}
|