mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(startup): 收口 leader 失锁后的后台任务
- 为后台调度器注册失锁回调并只在 stop 成功后清空生命周期引用 - 停止调度器时移除定时 job,补充启动与任务协调器回归测试 - 降低多 worker 下重复调度风险,保持停机收口与统计聚合回归一致
This commit is contained in:
89
src/main.py
89
src/main.py
@@ -7,6 +7,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import time
|
import time
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
@@ -114,6 +115,44 @@ class LifecycleState:
|
|||||||
warmup_task: asyncio.Task[None] | None = None
|
warmup_task: asyncio.Task[None] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
async def _stop_service_on_lock_lost(
|
||||||
|
state: LifecycleState,
|
||||||
|
*,
|
||||||
|
lock_name: str,
|
||||||
|
service_name: str,
|
||||||
|
state_attr: str,
|
||||||
|
stop: Callable[[], Awaitable[Any]],
|
||||||
|
) -> None:
|
||||||
|
logger.warning("检测到 {} 的 leader 锁已丢失,停止本实例的 {}", lock_name, service_name)
|
||||||
|
try:
|
||||||
|
await stop()
|
||||||
|
except Exception:
|
||||||
|
logger.exception("丢失 {} leader 锁后停止 {} 失败", lock_name, service_name)
|
||||||
|
return
|
||||||
|
|
||||||
|
setattr(state, state_attr, None)
|
||||||
|
|
||||||
|
|
||||||
|
def _make_lock_lost_callback(
|
||||||
|
state: LifecycleState,
|
||||||
|
*,
|
||||||
|
lock_name: str,
|
||||||
|
service_name: str,
|
||||||
|
state_attr: str,
|
||||||
|
stop: Callable[[], Awaitable[Any]],
|
||||||
|
) -> Callable[[str], Awaitable[None]]:
|
||||||
|
async def _callback(_name: str) -> None:
|
||||||
|
await _stop_service_on_lock_lost(
|
||||||
|
state,
|
||||||
|
lock_name=lock_name,
|
||||||
|
service_name=service_name,
|
||||||
|
state_attr=state_attr,
|
||||||
|
stop=stop,
|
||||||
|
)
|
||||||
|
|
||||||
|
return _callback
|
||||||
|
|
||||||
|
|
||||||
def _configure_uvicorn_access_log() -> None:
|
def _configure_uvicorn_access_log() -> None:
|
||||||
"""禁用 uvicorn access 日志(在子进程中执行)。"""
|
"""禁用 uvicorn access 日志(在子进程中执行)。"""
|
||||||
import logging
|
import logging
|
||||||
@@ -388,6 +427,16 @@ async def _start_background_services(state: LifecycleState) -> None:
|
|||||||
if quota_scheduler_active:
|
if quota_scheduler_active:
|
||||||
state.quota_scheduler = get_quota_scheduler()
|
state.quota_scheduler = get_quota_scheduler()
|
||||||
await state.quota_scheduler.start()
|
await state.quota_scheduler.start()
|
||||||
|
state.task_coordinator.register_lock_lost_callback(
|
||||||
|
"quota_scheduler",
|
||||||
|
_make_lock_lost_callback(
|
||||||
|
state,
|
||||||
|
lock_name="quota_scheduler",
|
||||||
|
service_name="月卡额度重置调度器",
|
||||||
|
state_attr="quota_scheduler",
|
||||||
|
stop=state.quota_scheduler.stop,
|
||||||
|
),
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行额度调度器,本实例跳过")
|
logger.info("检测到其他 worker 已运行额度调度器,本实例跳过")
|
||||||
state.quota_scheduler = None
|
state.quota_scheduler = None
|
||||||
@@ -398,6 +447,16 @@ async def _start_background_services(state: LifecycleState) -> None:
|
|||||||
state.maintenance_scheduler = get_maintenance_scheduler()
|
state.maintenance_scheduler = get_maintenance_scheduler()
|
||||||
logger.info("启动系统维护调度器...")
|
logger.info("启动系统维护调度器...")
|
||||||
await state.maintenance_scheduler.start()
|
await state.maintenance_scheduler.start()
|
||||||
|
state.task_coordinator.register_lock_lost_callback(
|
||||||
|
"maintenance_scheduler",
|
||||||
|
_make_lock_lost_callback(
|
||||||
|
state,
|
||||||
|
lock_name="maintenance_scheduler",
|
||||||
|
service_name="系统维护调度器",
|
||||||
|
state_attr="maintenance_scheduler",
|
||||||
|
stop=state.maintenance_scheduler.stop,
|
||||||
|
),
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行维护调度器,本实例跳过")
|
logger.info("检测到其他 worker 已运行维护调度器,本实例跳过")
|
||||||
state.maintenance_scheduler = None
|
state.maintenance_scheduler = None
|
||||||
@@ -408,6 +467,16 @@ async def _start_background_services(state: LifecycleState) -> None:
|
|||||||
state.model_fetch_scheduler = get_model_fetch_scheduler()
|
state.model_fetch_scheduler = get_model_fetch_scheduler()
|
||||||
logger.info("启动模型自动获取调度器...")
|
logger.info("启动模型自动获取调度器...")
|
||||||
await state.model_fetch_scheduler.start()
|
await state.model_fetch_scheduler.start()
|
||||||
|
state.task_coordinator.register_lock_lost_callback(
|
||||||
|
"model_fetch_scheduler",
|
||||||
|
_make_lock_lost_callback(
|
||||||
|
state,
|
||||||
|
lock_name="model_fetch_scheduler",
|
||||||
|
service_name="模型自动获取调度器",
|
||||||
|
state_attr="model_fetch_scheduler",
|
||||||
|
stop=state.model_fetch_scheduler.stop,
|
||||||
|
),
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行模型获取调度器,本实例跳过")
|
logger.info("检测到其他 worker 已运行模型获取调度器,本实例跳过")
|
||||||
state.model_fetch_scheduler = None
|
state.model_fetch_scheduler = None
|
||||||
@@ -420,6 +489,16 @@ async def _start_background_services(state: LifecycleState) -> None:
|
|||||||
state.pool_quota_probe_scheduler = get_pool_quota_probe_scheduler()
|
state.pool_quota_probe_scheduler = get_pool_quota_probe_scheduler()
|
||||||
logger.info("启动号池额度主动探测调度器...")
|
logger.info("启动号池额度主动探测调度器...")
|
||||||
await state.pool_quota_probe_scheduler.start()
|
await state.pool_quota_probe_scheduler.start()
|
||||||
|
state.task_coordinator.register_lock_lost_callback(
|
||||||
|
"pool_quota_probe_scheduler",
|
||||||
|
_make_lock_lost_callback(
|
||||||
|
state,
|
||||||
|
lock_name="pool_quota_probe_scheduler",
|
||||||
|
service_name="号池额度主动探测调度器",
|
||||||
|
state_attr="pool_quota_probe_scheduler",
|
||||||
|
stop=state.pool_quota_probe_scheduler.stop,
|
||||||
|
),
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行号池额度主动探测调度器,本实例跳过")
|
logger.info("检测到其他 worker 已运行号池额度主动探测调度器,本实例跳过")
|
||||||
state.pool_quota_probe_scheduler = None
|
state.pool_quota_probe_scheduler = None
|
||||||
@@ -430,6 +509,16 @@ async def _start_background_services(state: LifecycleState) -> None:
|
|||||||
state.task_poller = get_task_poller()
|
state.task_poller = get_task_poller()
|
||||||
logger.info("启动 TaskPoller(video)...")
|
logger.info("启动 TaskPoller(video)...")
|
||||||
await state.task_poller.start()
|
await state.task_poller.start()
|
||||||
|
state.task_coordinator.register_lock_lost_callback(
|
||||||
|
"task_poller:video",
|
||||||
|
_make_lock_lost_callback(
|
||||||
|
state,
|
||||||
|
lock_name="task_poller:video",
|
||||||
|
service_name="TaskPoller(video)",
|
||||||
|
state_attr="task_poller",
|
||||||
|
stop=state.task_poller.stop,
|
||||||
|
),
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行 TaskPoller(video),本实例跳过")
|
logger.info("检测到其他 worker 已运行 TaskPoller(video),本实例跳过")
|
||||||
state.task_poller = None
|
state.task_poller = None
|
||||||
|
|||||||
@@ -20,6 +20,9 @@ if TYPE_CHECKING:
|
|||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
|
|
||||||
|
_proxy_node_task_coordinator: Any | None = None
|
||||||
|
|
||||||
|
|
||||||
def _reset_tunnel_connected_on_startup() -> None:
|
def _reset_tunnel_connected_on_startup() -> None:
|
||||||
"""服务端启动时将所有 tunnel_connected=True 的节点重置为 False/OFFLINE。
|
"""服务端启动时将所有 tunnel_connected=True 的节点重置为 False/OFFLINE。
|
||||||
|
|
||||||
@@ -89,8 +92,10 @@ async def _on_startup() -> None:
|
|||||||
|
|
||||||
from src.clients import get_redis_client
|
from src.clients import get_redis_client
|
||||||
|
|
||||||
|
global _proxy_node_task_coordinator
|
||||||
redis_client = await get_redis_client()
|
redis_client = await get_redis_client()
|
||||||
task_coordinator = StartupTaskCoordinator(redis_client)
|
task_coordinator = StartupTaskCoordinator(redis_client)
|
||||||
|
_proxy_node_task_coordinator = task_coordinator
|
||||||
|
|
||||||
proxy_node_health_scheduler = get_proxy_node_health_scheduler()
|
proxy_node_health_scheduler = get_proxy_node_health_scheduler()
|
||||||
active = await task_coordinator.acquire("proxy_node_health")
|
active = await task_coordinator.acquire("proxy_node_health")
|
||||||
@@ -104,6 +109,10 @@ async def _on_startup() -> None:
|
|||||||
if active:
|
if active:
|
||||||
logger.info("启动 ProxyNode 心跳检测调度器...")
|
logger.info("启动 ProxyNode 心跳检测调度器...")
|
||||||
await proxy_node_health_scheduler.start()
|
await proxy_node_health_scheduler.start()
|
||||||
|
task_coordinator.register_lock_lost_callback(
|
||||||
|
"proxy_node_health",
|
||||||
|
lambda _name: _stop_proxy_node_scheduler_on_lock_lost(logger),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def _on_shutdown() -> None:
|
async def _on_shutdown() -> None:
|
||||||
@@ -117,14 +126,28 @@ async def _on_shutdown() -> None:
|
|||||||
|
|
||||||
from src.clients import get_redis_client
|
from src.clients import get_redis_client
|
||||||
|
|
||||||
|
global _proxy_node_task_coordinator
|
||||||
|
if _proxy_node_task_coordinator is None:
|
||||||
redis_client = await get_redis_client()
|
redis_client = await get_redis_client()
|
||||||
task_coordinator = StartupTaskCoordinator(redis_client)
|
_proxy_node_task_coordinator = StartupTaskCoordinator(redis_client)
|
||||||
|
|
||||||
scheduler = get_proxy_node_health_scheduler()
|
scheduler = get_proxy_node_health_scheduler()
|
||||||
if scheduler.running:
|
if scheduler.running:
|
||||||
logger.info("停止 ProxyNode 心跳检测调度器...")
|
logger.info("停止 ProxyNode 心跳检测调度器...")
|
||||||
await scheduler.stop()
|
await scheduler.stop()
|
||||||
await task_coordinator.release("proxy_node_health")
|
await _proxy_node_task_coordinator.release("proxy_node_health")
|
||||||
|
_proxy_node_task_coordinator = None
|
||||||
|
|
||||||
|
|
||||||
|
async def _stop_proxy_node_scheduler_on_lock_lost(logger: Any) -> None:
|
||||||
|
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||||
|
|
||||||
|
scheduler = get_proxy_node_health_scheduler()
|
||||||
|
if not scheduler.running:
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.warning("检测到 proxy_node_health 的 leader 锁已丢失,停止本实例的 ProxyNode 心跳检测")
|
||||||
|
await scheduler.stop()
|
||||||
|
|
||||||
|
|
||||||
async def _health_check() -> ModuleHealth:
|
async def _health_check() -> ModuleHealth:
|
||||||
|
|||||||
@@ -54,9 +54,7 @@ KEY_FETCH_TIMEOUT_SECONDS = 120
|
|||||||
MODEL_FETCH_HTTP_TIMEOUT = 10.0
|
MODEL_FETCH_HTTP_TIMEOUT = 10.0
|
||||||
|
|
||||||
# 启动时首次自动获取开关与延迟
|
# 启动时首次自动获取开关与延迟
|
||||||
MODEL_FETCH_STARTUP_ENABLED = (
|
MODEL_FETCH_STARTUP_ENABLED = os.getenv("MODEL_FETCH_STARTUP_ENABLED", "true").lower() == "true"
|
||||||
os.getenv("MODEL_FETCH_STARTUP_ENABLED", "true").lower() == "true"
|
|
||||||
)
|
|
||||||
MODEL_FETCH_STARTUP_DELAY_SECONDS = max(
|
MODEL_FETCH_STARTUP_DELAY_SECONDS = max(
|
||||||
0,
|
0,
|
||||||
int(os.getenv("MODEL_FETCH_STARTUP_DELAY_SECONDS", "10")),
|
int(os.getenv("MODEL_FETCH_STARTUP_DELAY_SECONDS", "10")),
|
||||||
@@ -315,6 +313,8 @@ class ModelFetchScheduler:
|
|||||||
async def stop(self) -> None:
|
async def stop(self) -> None:
|
||||||
"""停止调度器"""
|
"""停止调度器"""
|
||||||
self._running = False
|
self._running = False
|
||||||
|
scheduler = get_scheduler()
|
||||||
|
scheduler.remove_job("model_auto_fetch")
|
||||||
|
|
||||||
# 取消并等待启动任务完成
|
# 取消并等待启动任务完成
|
||||||
if self._startup_task and not self._startup_task.done():
|
if self._startup_task and not self._startup_task.done():
|
||||||
@@ -348,6 +348,8 @@ class ModelFetchScheduler:
|
|||||||
|
|
||||||
async def _scheduled_fetch_models(self) -> None:
|
async def _scheduled_fetch_models(self) -> None:
|
||||||
"""定时任务入口"""
|
"""定时任务入口"""
|
||||||
|
if not self._running:
|
||||||
|
return
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
await self._perform_fetch_all_keys()
|
await self._perform_fetch_all_keys()
|
||||||
|
|
||||||
|
|||||||
@@ -183,6 +183,8 @@ class PoolQuotaProbeScheduler:
|
|||||||
if not self.running:
|
if not self.running:
|
||||||
return
|
return
|
||||||
self.running = False
|
self.running = False
|
||||||
|
scheduler = get_scheduler()
|
||||||
|
scheduler.remove_job("pool_quota_probe_check")
|
||||||
logger.info("PoolQuotaProbeScheduler stopped")
|
logger.info("PoolQuotaProbeScheduler stopped")
|
||||||
|
|
||||||
async def _scheduled_probe_check(self) -> None:
|
async def _scheduled_probe_check(self) -> None:
|
||||||
|
|||||||
@@ -74,9 +74,13 @@ class ProxyNodeHealthScheduler:
|
|||||||
if not self.running:
|
if not self.running:
|
||||||
return
|
return
|
||||||
self.running = False
|
self.running = False
|
||||||
|
scheduler = get_scheduler()
|
||||||
|
scheduler.remove_job("proxy_node_health_check")
|
||||||
logger.info("ProxyNodeHealthScheduler stopped")
|
logger.info("ProxyNodeHealthScheduler stopped")
|
||||||
|
|
||||||
async def _scheduled_check(self) -> None:
|
async def _scheduled_check(self) -> None:
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._check_heartbeats()
|
await self._check_heartbeats()
|
||||||
self._check_count = (self._check_count + 1) % _EVENT_CLEANUP_INTERVAL
|
self._check_count = (self._check_count + 1) % _EVENT_CLEANUP_INTERVAL
|
||||||
if self._check_count == 0:
|
if self._check_count == 0:
|
||||||
|
|||||||
@@ -49,6 +49,7 @@ class MaintenanceScheduler:
|
|||||||
self._interval_tasks = []
|
self._interval_tasks = []
|
||||||
self._stats_aggregation_lock = asyncio.Lock()
|
self._stats_aggregation_lock = asyncio.Lock()
|
||||||
self._wallet_daily_usage_lock = asyncio.Lock()
|
self._wallet_daily_usage_lock = asyncio.Lock()
|
||||||
|
self._startup_task: asyncio.Task[Any] | None = None
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _get_http_client_idle_cleanup_interval_minutes() -> int:
|
def _get_http_client_idle_cleanup_interval_minutes() -> int:
|
||||||
@@ -260,8 +261,9 @@ class MaintenanceScheduler:
|
|||||||
if config.maintenance_startup_tasks_enabled:
|
if config.maintenance_startup_tasks_enabled:
|
||||||
from src.utils.async_utils import safe_create_task
|
from src.utils.async_utils import safe_create_task
|
||||||
|
|
||||||
safe_create_task(self._run_startup_tasks())
|
self._startup_task = safe_create_task(self._run_startup_tasks())
|
||||||
else:
|
else:
|
||||||
|
self._startup_task = None
|
||||||
logger.info("维护调度器启动任务已禁用(MAINTENANCE_STARTUP_TASKS_ENABLED=false)")
|
logger.info("维护调度器启动任务已禁用(MAINTENANCE_STARTUP_TASKS_ENABLED=false)")
|
||||||
|
|
||||||
async def _run_startup_tasks(self) -> None:
|
async def _run_startup_tasks(self) -> None:
|
||||||
@@ -290,8 +292,32 @@ class MaintenanceScheduler:
|
|||||||
return
|
return
|
||||||
|
|
||||||
self.running = False
|
self.running = False
|
||||||
|
|
||||||
|
if self._startup_task and not self._startup_task.done():
|
||||||
|
self._startup_task.cancel()
|
||||||
|
try:
|
||||||
|
await self._startup_task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
self._startup_task = None
|
||||||
|
|
||||||
scheduler = get_scheduler()
|
scheduler = get_scheduler()
|
||||||
scheduler.stop()
|
for job_id in (
|
||||||
|
"stats_aggregation",
|
||||||
|
"stats_hourly_aggregation",
|
||||||
|
"wallet_daily_usage_aggregation",
|
||||||
|
"usage_cleanup",
|
||||||
|
"pool_monitor",
|
||||||
|
"http_client_idle_cleanup",
|
||||||
|
"pending_cleanup",
|
||||||
|
"audit_cleanup",
|
||||||
|
"gemini_file_mapping_cleanup",
|
||||||
|
"candidate_cleanup",
|
||||||
|
"db_maintenance",
|
||||||
|
"antigravity_ua_refresh",
|
||||||
|
self.CHECKIN_JOB_ID,
|
||||||
|
):
|
||||||
|
scheduler.remove_job(job_id)
|
||||||
|
|
||||||
logger.info("系统维护调度器已停止")
|
logger.info("系统维护调度器已停止")
|
||||||
|
|
||||||
@@ -299,22 +325,32 @@ class MaintenanceScheduler:
|
|||||||
|
|
||||||
async def _scheduled_stats_aggregation(self, backfill: bool = False) -> None:
|
async def _scheduled_stats_aggregation(self, backfill: bool = False) -> None:
|
||||||
"""统计聚合任务(定时调用)"""
|
"""统计聚合任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_stats_aggregation(backfill=backfill)
|
await self._perform_stats_aggregation(backfill=backfill)
|
||||||
|
|
||||||
async def _scheduled_wallet_daily_usage_aggregation(self) -> None:
|
async def _scheduled_wallet_daily_usage_aggregation(self) -> None:
|
||||||
"""钱包每日消费汇总任务(定时调用)"""
|
"""钱包每日消费汇总任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_wallet_daily_usage_aggregation()
|
await self._perform_wallet_daily_usage_aggregation()
|
||||||
|
|
||||||
async def _scheduled_hourly_stats_aggregation(self) -> None:
|
async def _scheduled_hourly_stats_aggregation(self) -> None:
|
||||||
"""小时统计聚合任务(定时调用)"""
|
"""小时统计聚合任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_hourly_stats_aggregation()
|
await self._perform_hourly_stats_aggregation()
|
||||||
|
|
||||||
async def _scheduled_cleanup(self) -> None:
|
async def _scheduled_cleanup(self) -> None:
|
||||||
"""清理任务(定时调用)"""
|
"""清理任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_cleanup()
|
await self._perform_cleanup()
|
||||||
|
|
||||||
async def _scheduled_monitor(self) -> None:
|
async def _scheduled_monitor(self) -> None:
|
||||||
"""监控任务(定时调用)"""
|
"""监控任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
try:
|
try:
|
||||||
from src.database import log_pool_status
|
from src.database import log_pool_status
|
||||||
|
|
||||||
@@ -324,6 +360,8 @@ class MaintenanceScheduler:
|
|||||||
|
|
||||||
async def _scheduled_http_client_idle_cleanup(self) -> None:
|
async def _scheduled_http_client_idle_cleanup(self) -> None:
|
||||||
"""HTTP 客户端空闲清理任务(定时调用)。"""
|
"""HTTP 客户端空闲清理任务(定时调用)。"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
try:
|
try:
|
||||||
stats = await HTTPClientPool.cleanup_idle_clients()
|
stats = await HTTPClientPool.cleanup_idle_clients()
|
||||||
if stats.get("proxy_closed", 0) or stats.get("tunnel_closed", 0):
|
if stats.get("proxy_closed", 0) or stats.get("tunnel_closed", 0):
|
||||||
@@ -337,26 +375,38 @@ class MaintenanceScheduler:
|
|||||||
|
|
||||||
async def _scheduled_pending_cleanup(self) -> None:
|
async def _scheduled_pending_cleanup(self) -> None:
|
||||||
"""Pending 清理任务(定时调用)"""
|
"""Pending 清理任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_pending_cleanup()
|
await self._perform_pending_cleanup()
|
||||||
|
|
||||||
async def _scheduled_audit_cleanup(self) -> None:
|
async def _scheduled_audit_cleanup(self) -> None:
|
||||||
"""审计日志清理任务(定时调用)"""
|
"""审计日志清理任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_audit_cleanup()
|
await self._perform_audit_cleanup()
|
||||||
|
|
||||||
async def _scheduled_candidate_cleanup(self) -> None:
|
async def _scheduled_candidate_cleanup(self) -> None:
|
||||||
"""请求候选记录清理任务(定时调用)"""
|
"""请求候选记录清理任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_candidate_cleanup()
|
await self._perform_candidate_cleanup()
|
||||||
|
|
||||||
async def _scheduled_db_maintenance(self) -> None:
|
async def _scheduled_db_maintenance(self) -> None:
|
||||||
"""数据库表维护任务(定时调用)"""
|
"""数据库表维护任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_db_maintenance()
|
await self._perform_db_maintenance()
|
||||||
|
|
||||||
async def _scheduled_gemini_file_mapping_cleanup(self) -> None:
|
async def _scheduled_gemini_file_mapping_cleanup(self) -> None:
|
||||||
"""Gemini 文件映射清理任务(定时调用)"""
|
"""Gemini 文件映射清理任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_gemini_file_mapping_cleanup()
|
await self._perform_gemini_file_mapping_cleanup()
|
||||||
|
|
||||||
async def _scheduled_antigravity_ua_refresh(self) -> None:
|
async def _scheduled_antigravity_ua_refresh(self) -> None:
|
||||||
"""Antigravity User-Agent 版本刷新(定时调用)"""
|
"""Antigravity User-Agent 版本刷新(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
try:
|
try:
|
||||||
from src.services.provider.adapters.antigravity.client import refresh_user_agent
|
from src.services.provider.adapters.antigravity.client import refresh_user_agent
|
||||||
|
|
||||||
@@ -366,6 +416,8 @@ class MaintenanceScheduler:
|
|||||||
|
|
||||||
async def _scheduled_provider_checkin(self) -> None:
|
async def _scheduled_provider_checkin(self) -> None:
|
||||||
"""Provider 签到任务(定时调用)"""
|
"""Provider 签到任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._perform_provider_checkin()
|
await self._perform_provider_checkin()
|
||||||
|
|
||||||
# ========== 实际任务实现 ==========
|
# ========== 实际任务实现 ==========
|
||||||
|
|||||||
@@ -873,6 +873,8 @@ class StatsAggregatorService:
|
|||||||
db: Session, date: datetime, user_ids: list[str] | None = None
|
db: Session, date: datetime, user_ids: list[str] | None = None
|
||||||
) -> StatsDaily:
|
) -> StatsDaily:
|
||||||
"""聚合单日所有统计(原子提交)"""
|
"""聚合单日所有统计(原子提交)"""
|
||||||
|
|
||||||
|
def _do_aggregate() -> StatsDaily:
|
||||||
stats = StatsAggregatorService.aggregate_daily_stats(db, date, commit=False)
|
stats = StatsAggregatorService.aggregate_daily_stats(db, date, commit=False)
|
||||||
StatsAggregatorService.aggregate_daily_model_stats(db, date, commit=False)
|
StatsAggregatorService.aggregate_daily_model_stats(db, date, commit=False)
|
||||||
StatsAggregatorService.aggregate_daily_provider_stats(db, date, commit=False)
|
StatsAggregatorService.aggregate_daily_provider_stats(db, date, commit=False)
|
||||||
@@ -889,6 +891,13 @@ class StatsAggregatorService:
|
|||||||
db.commit()
|
db.commit()
|
||||||
return stats
|
return stats
|
||||||
|
|
||||||
|
try:
|
||||||
|
return _do_aggregate()
|
||||||
|
except IntegrityError:
|
||||||
|
db.rollback()
|
||||||
|
logger.warning("每日统计聚合冲突,重试更新: {}", date)
|
||||||
|
return _do_aggregate()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def aggregate_hourly_stats(db: Session, hour_utc: datetime, commit: bool = True) -> StatsHourly:
|
def aggregate_hourly_stats(db: Session, hour_utc: datetime, commit: bool = True) -> StatsHourly:
|
||||||
"""聚合指定 UTC 小时的全局统计"""
|
"""聚合指定 UTC 小时的全局统计"""
|
||||||
|
|||||||
@@ -54,10 +54,14 @@ class QuotaScheduler:
|
|||||||
return
|
return
|
||||||
|
|
||||||
self.running = False
|
self.running = False
|
||||||
|
scheduler = get_scheduler()
|
||||||
|
scheduler.remove_job("quota_reset_check")
|
||||||
logger.info("Quota scheduler stopped")
|
logger.info("Quota scheduler stopped")
|
||||||
|
|
||||||
async def _scheduled_quota_check(self) -> None:
|
async def _scheduled_quota_check(self) -> None:
|
||||||
"""额度检查任务(定时调用)"""
|
"""额度检查任务(定时调用)"""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
await self._check_and_reset_quotas()
|
await self._check_and_reset_quotas()
|
||||||
|
|
||||||
async def _check_and_reset_quotas(self) -> None:
|
async def _check_and_reset_quotas(self) -> None:
|
||||||
|
|||||||
@@ -11,9 +11,11 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import pathlib
|
import pathlib
|
||||||
import uuid
|
import uuid
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
@@ -35,6 +37,8 @@ class StartupTaskCoordinator:
|
|||||||
self.redis = redis_client
|
self.redis = redis_client
|
||||||
self._tokens: dict[str, str] = {}
|
self._tokens: dict[str, str] = {}
|
||||||
self._file_handles: dict[str, object] = {}
|
self._file_handles: dict[str, object] = {}
|
||||||
|
self._refresh_tasks: dict[str, asyncio.Task[None]] = {}
|
||||||
|
self._lock_lost_callbacks: dict[str, Callable[[str], Awaitable[None] | None]] = {}
|
||||||
self._lock_dir = pathlib.Path(lock_dir or os.getenv("TASK_LOCK_DIR", "./.locks"))
|
self._lock_dir = pathlib.Path(lock_dir or os.getenv("TASK_LOCK_DIR", "./.locks"))
|
||||||
if not self._lock_dir.exists():
|
if not self._lock_dir.exists():
|
||||||
self._lock_dir.mkdir(parents=True, exist_ok=True)
|
self._lock_dir.mkdir(parents=True, exist_ok=True)
|
||||||
@@ -80,6 +84,7 @@ class StartupTaskCoordinator:
|
|||||||
)
|
)
|
||||||
if result == 1:
|
if result == 1:
|
||||||
self._tokens[name] = token
|
self._tokens[name] = token
|
||||||
|
self._start_refresh_task(name, ttl)
|
||||||
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
@@ -88,6 +93,7 @@ class StartupTaskCoordinator:
|
|||||||
acquired = await self.redis.set(self._redis_key(name), token, nx=True, ex=ttl)
|
acquired = await self.redis.set(self._redis_key(name), token, nx=True, ex=ttl)
|
||||||
if acquired:
|
if acquired:
|
||||||
self._tokens[name] = token
|
self._tokens[name] = token
|
||||||
|
self._start_refresh_task(name, ttl)
|
||||||
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
@@ -97,6 +103,11 @@ class StartupTaskCoordinator:
|
|||||||
return await self._acquire_file_lock(name)
|
return await self._acquire_file_lock(name)
|
||||||
|
|
||||||
async def release(self, name: str) -> Any:
|
async def release(self, name: str) -> Any:
|
||||||
|
refresh_task = self._refresh_tasks.pop(name, None)
|
||||||
|
if refresh_task is not None:
|
||||||
|
refresh_task.cancel()
|
||||||
|
self._lock_lost_callbacks.pop(name, None)
|
||||||
|
|
||||||
if self.redis and name in self._tokens:
|
if self.redis and name in self._tokens:
|
||||||
token = self._tokens.pop(name)
|
token = self._tokens.pop(name)
|
||||||
script = """
|
script = """
|
||||||
@@ -137,6 +148,72 @@ class StartupTaskCoordinator:
|
|||||||
handle.close()
|
handle.close()
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def register_lock_lost_callback(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
callback: Callable[[str], Awaitable[None] | None],
|
||||||
|
) -> None:
|
||||||
|
self._lock_lost_callbacks[name] = callback
|
||||||
|
|
||||||
|
def _start_refresh_task(self, name: str, ttl: int) -> None:
|
||||||
|
if not self.redis:
|
||||||
|
return
|
||||||
|
|
||||||
|
existing_task = self._refresh_tasks.pop(name, None)
|
||||||
|
if existing_task is not None:
|
||||||
|
existing_task.cancel()
|
||||||
|
|
||||||
|
self._refresh_tasks[name] = asyncio.create_task(self._refresh_lock_loop(name, ttl))
|
||||||
|
|
||||||
|
async def _refresh_lock(self, name: str, ttl: int) -> bool:
|
||||||
|
token = self._tokens.get(name)
|
||||||
|
if not self.redis or token is None:
|
||||||
|
return False
|
||||||
|
|
||||||
|
script = """
|
||||||
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||||||
|
return redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]))
|
||||||
|
end
|
||||||
|
return 0
|
||||||
|
"""
|
||||||
|
result = await self.redis.eval(script, 1, self._redis_key(name), token, ttl)
|
||||||
|
return result == 1
|
||||||
|
|
||||||
|
async def _refresh_lock_loop(self, name: str, ttl: int) -> None:
|
||||||
|
interval = max(1, ttl // 3)
|
||||||
|
try:
|
||||||
|
while name in self._tokens:
|
||||||
|
await asyncio.sleep(interval)
|
||||||
|
if name not in self._tokens:
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
refreshed = await self._refresh_lock(name, ttl)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception as exc: # pragma: no cover - 续租失败后下轮重试
|
||||||
|
logger.warning(f"续租任务锁 {name} 失败: {exc}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
if not refreshed:
|
||||||
|
logger.warning(f"任务 {name} 的 Redis 锁已失效,停止续租")
|
||||||
|
self._tokens.pop(name, None)
|
||||||
|
await self._notify_lock_lost(name)
|
||||||
|
break
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
|
||||||
|
async def _notify_lock_lost(self, name: str) -> None:
|
||||||
|
callback = self._lock_lost_callbacks.pop(name, None)
|
||||||
|
if callback is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = callback(name)
|
||||||
|
if result is not None:
|
||||||
|
await result
|
||||||
|
except Exception as exc: # pragma: no cover - 回调失败仅记录日志
|
||||||
|
logger.exception("任务 {} 的失锁回调执行失败: {}", name, exc)
|
||||||
|
|
||||||
|
|
||||||
async def ensure_singleton_task(
|
async def ensure_singleton_task(
|
||||||
name: str, redis_client: Any | None = None, ttl: int | None = None
|
name: str, redis_client: Any | None = None, ttl: int | None = None
|
||||||
|
|||||||
@@ -1,12 +1,16 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import inspect
|
import inspect
|
||||||
|
from collections.abc import Callable
|
||||||
from datetime import datetime, timedelta, timezone
|
from datetime import datetime, timedelta, timezone
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
|
from typing import Any, cast
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
import src.main as main_module
|
||||||
import src.services.system.maintenance_scheduler as maintenance_scheduler_module
|
import src.services.system.maintenance_scheduler as maintenance_scheduler_module
|
||||||
from src.config.settings import config
|
from src.config.settings import config
|
||||||
from src.services.system.maintenance_scheduler import MaintenanceScheduler
|
from src.services.system.maintenance_scheduler import MaintenanceScheduler
|
||||||
@@ -45,6 +49,85 @@ async def test_maintenance_scheduler_start_skips_startup_task_when_disabled(
|
|||||||
assert created is False
|
assert created is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_maintenance_scheduler_stop_cancels_startup_task_and_removes_jobs(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
scheduler = MaintenanceScheduler()
|
||||||
|
scheduler.running = True
|
||||||
|
scheduler._startup_task = asyncio.create_task(asyncio.sleep(3600))
|
||||||
|
removed_jobs: list[str] = []
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
maintenance_scheduler_module,
|
||||||
|
"get_scheduler",
|
||||||
|
lambda: SimpleNamespace(remove_job=lambda job_id: removed_jobs.append(job_id)),
|
||||||
|
)
|
||||||
|
|
||||||
|
await scheduler.stop()
|
||||||
|
|
||||||
|
assert scheduler.running is False
|
||||||
|
assert scheduler._startup_task is None
|
||||||
|
assert set(removed_jobs) == {
|
||||||
|
"stats_aggregation",
|
||||||
|
"stats_hourly_aggregation",
|
||||||
|
"wallet_daily_usage_aggregation",
|
||||||
|
"usage_cleanup",
|
||||||
|
"pool_monitor",
|
||||||
|
"http_client_idle_cleanup",
|
||||||
|
"pending_cleanup",
|
||||||
|
"audit_cleanup",
|
||||||
|
"gemini_file_mapping_cleanup",
|
||||||
|
"candidate_cleanup",
|
||||||
|
"db_maintenance",
|
||||||
|
"antigravity_ua_refresh",
|
||||||
|
scheduler.CHECKIN_JOB_ID,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stop_service_on_lock_lost_keeps_state_when_stop_fails() -> None:
|
||||||
|
state = main_module.LifecycleState()
|
||||||
|
service = SimpleNamespace()
|
||||||
|
state.quota_scheduler = cast(Any, service)
|
||||||
|
|
||||||
|
async def fail_stop() -> None:
|
||||||
|
raise RuntimeError("boom")
|
||||||
|
|
||||||
|
await main_module._stop_service_on_lock_lost(
|
||||||
|
state,
|
||||||
|
lock_name="quota_scheduler",
|
||||||
|
service_name="月卡额度重置调度器",
|
||||||
|
state_attr="quota_scheduler",
|
||||||
|
stop=fail_stop,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert state.quota_scheduler is service
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stop_service_on_lock_lost_clears_state_after_success() -> None:
|
||||||
|
state = main_module.LifecycleState()
|
||||||
|
service = SimpleNamespace()
|
||||||
|
state.quota_scheduler = cast(Any, service)
|
||||||
|
stopped = False
|
||||||
|
|
||||||
|
async def stop() -> None:
|
||||||
|
nonlocal stopped
|
||||||
|
stopped = True
|
||||||
|
|
||||||
|
await main_module._stop_service_on_lock_lost(
|
||||||
|
state,
|
||||||
|
lock_name="quota_scheduler",
|
||||||
|
service_name="月卡额度重置调度器",
|
||||||
|
state_attr="quota_scheduler",
|
||||||
|
stop=stop,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert stopped is True
|
||||||
|
assert state.quota_scheduler is None
|
||||||
|
|
||||||
|
|
||||||
def test_http_client_idle_cleanup_interval_env_invalid(
|
def test_http_client_idle_cleanup_interval_env_invalid(
|
||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -276,7 +359,7 @@ async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
|
|||||||
|
|
||||||
class _FakeDateTime(datetime):
|
class _FakeDateTime(datetime):
|
||||||
@classmethod
|
@classmethod
|
||||||
def now(cls, tz=None): # type: ignore[override]
|
def now(cls, tz: timezone | None = None) -> datetime: # type: ignore[override]
|
||||||
if tz is None:
|
if tz is None:
|
||||||
return fixed_now.replace(tzinfo=None)
|
return fixed_now.replace(tzinfo=None)
|
||||||
return fixed_now.astimezone(tz)
|
return fixed_now.astimezone(tz)
|
||||||
@@ -291,7 +374,7 @@ async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
|
|||||||
|
|
||||||
calls: list[tuple[str, datetime, int, datetime | None]] = []
|
calls: list[tuple[str, datetime, int, datetime | None]] = []
|
||||||
|
|
||||||
def _record(name: str, count: int):
|
def _record(name: str, count: int) -> Callable[..., int]:
|
||||||
def _inner(
|
def _inner(
|
||||||
cutoff_time: datetime,
|
cutoff_time: datetime,
|
||||||
batch_size: int,
|
batch_size: int,
|
||||||
@@ -313,6 +396,10 @@ async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
|
|||||||
"auto_delete_expired_keys": False,
|
"auto_delete_expired_keys": False,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
def _delete_old_records(cutoff_time: datetime, batch_size: int) -> int:
|
||||||
|
calls.append(("delete", cutoff_time, batch_size, None))
|
||||||
|
return 5
|
||||||
|
|
||||||
monkeypatch.setattr(maintenance_scheduler_module, "datetime", _FakeDateTime)
|
monkeypatch.setattr(maintenance_scheduler_module, "datetime", _FakeDateTime)
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
|
maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
|
||||||
@@ -330,8 +417,7 @@ async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
|
|||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
scheduler,
|
scheduler,
|
||||||
"_delete_old_records",
|
"_delete_old_records",
|
||||||
lambda cutoff_time, batch_size: calls.append(("delete", cutoff_time, batch_size, None))
|
_delete_old_records,
|
||||||
or 5,
|
|
||||||
)
|
)
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
scheduler,
|
scheduler,
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ from types import SimpleNamespace
|
|||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
|
||||||
from src.models.database import StatsDaily, StatsUserDaily
|
from src.models.database import StatsDaily, StatsUserDaily
|
||||||
from src.services.system.stats_aggregator import (
|
from src.services.system.stats_aggregator import (
|
||||||
@@ -61,6 +62,20 @@ class _BatchUserStatsSession:
|
|||||||
self.commit_count += 1
|
self.commit_count += 1
|
||||||
|
|
||||||
|
|
||||||
|
class _RetryCommitSession:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.commit_count = 0
|
||||||
|
self.rollback_count = 0
|
||||||
|
|
||||||
|
def commit(self) -> None:
|
||||||
|
self.commit_count += 1
|
||||||
|
if self.commit_count == 1:
|
||||||
|
raise IntegrityError("insert", {}, Exception("duplicate key"))
|
||||||
|
|
||||||
|
def rollback(self) -> None:
|
||||||
|
self.rollback_count += 1
|
||||||
|
|
||||||
|
|
||||||
def test_query_stats_hybrid_batches_statsdaily_lookup_and_merges_realtime_ranges(
|
def test_query_stats_hybrid_batches_statsdaily_lookup_and_merges_realtime_ranges(
|
||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -167,6 +182,79 @@ def test_aggregate_user_daily_stats_batch_updates_all_users_in_two_queries() ->
|
|||||||
assert user_two.total_cost == 0.0
|
assert user_two.total_cost == 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_aggregate_daily_stats_bundle_retries_after_integrity_error(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
target_day = datetime(2026, 3, 1, tzinfo=timezone.utc)
|
||||||
|
db = _RetryCommitSession()
|
||||||
|
stats_calls: list[SimpleNamespace] = []
|
||||||
|
stage_calls: list[str] = []
|
||||||
|
|
||||||
|
def fake_aggregate_daily_stats(*_args: object, **_kwargs: object) -> SimpleNamespace:
|
||||||
|
stats = SimpleNamespace(is_complete=False, aggregated_at=None)
|
||||||
|
stats_calls.append(stats)
|
||||||
|
stage_calls.append("daily")
|
||||||
|
return stats
|
||||||
|
|
||||||
|
def fake_model_stats(*_args: object, **_kwargs: object) -> list[object]:
|
||||||
|
stage_calls.append("model")
|
||||||
|
return []
|
||||||
|
|
||||||
|
def fake_provider_stats(*_args: object, **_kwargs: object) -> list[object]:
|
||||||
|
stage_calls.append("provider")
|
||||||
|
return []
|
||||||
|
|
||||||
|
def fake_api_key_stats(*_args: object, **_kwargs: object) -> list[object]:
|
||||||
|
stage_calls.append("api_key")
|
||||||
|
return []
|
||||||
|
|
||||||
|
def fake_error_stats(*_args: object, **_kwargs: object) -> list[object]:
|
||||||
|
stage_calls.append("error")
|
||||||
|
return []
|
||||||
|
|
||||||
|
def fake_user_daily_stats(*_args: object, **_kwargs: object) -> list[object]:
|
||||||
|
stage_calls.append("user")
|
||||||
|
return []
|
||||||
|
|
||||||
|
monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_stats", fake_aggregate_daily_stats)
|
||||||
|
monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_model_stats", fake_model_stats)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
StatsAggregatorService, "aggregate_daily_provider_stats", fake_provider_stats
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_api_key_stats", fake_api_key_stats)
|
||||||
|
monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_error_stats", fake_error_stats)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
StatsAggregatorService, "aggregate_user_daily_stats_batch", fake_user_daily_stats
|
||||||
|
)
|
||||||
|
|
||||||
|
result = StatsAggregatorService.aggregate_daily_stats_bundle(
|
||||||
|
cast(Any, db),
|
||||||
|
target_day,
|
||||||
|
user_ids=["user-1"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert db.commit_count == 2
|
||||||
|
assert db.rollback_count == 1
|
||||||
|
assert len(stats_calls) == 2
|
||||||
|
assert stage_calls == [
|
||||||
|
"daily",
|
||||||
|
"model",
|
||||||
|
"provider",
|
||||||
|
"api_key",
|
||||||
|
"error",
|
||||||
|
"user",
|
||||||
|
"daily",
|
||||||
|
"model",
|
||||||
|
"provider",
|
||||||
|
"api_key",
|
||||||
|
"error",
|
||||||
|
"user",
|
||||||
|
]
|
||||||
|
assert result is stats_calls[-1]
|
||||||
|
assert result.is_complete is True
|
||||||
|
assert result.aggregated_at is not None
|
||||||
|
|
||||||
|
|
||||||
def test_compute_percentiles_by_local_day_returns_sqlite_fallback_without_queries() -> None:
|
def test_compute_percentiles_by_local_day_returns_sqlite_fallback_without_queries() -> None:
|
||||||
db = SimpleNamespace(bind=SimpleNamespace(dialect=SimpleNamespace(name="sqlite")))
|
db = SimpleNamespace(bind=SimpleNamespace(dialect=SimpleNamespace(name="sqlite")))
|
||||||
time_range = TimeRangeParams(
|
time_range = TimeRangeParams(
|
||||||
|
|||||||
106
tests/unit/test_task_coordinator.py
Normal file
106
tests/unit/test_task_coordinator.py
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.utils.task_coordinator import StartupTaskCoordinator
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeRedis:
|
||||||
|
def __init__(self, eval_results: list[int] | None = None, set_result: bool = True) -> None:
|
||||||
|
self.eval_results = list(eval_results or [])
|
||||||
|
self.eval_calls: list[tuple[Any, ...]] = []
|
||||||
|
self.set_result = set_result
|
||||||
|
self.set_calls: list[tuple[Any, ...]] = []
|
||||||
|
|
||||||
|
async def eval(self, script: str, numkeys: int, *args: Any) -> int:
|
||||||
|
self.eval_calls.append((script, numkeys, *args))
|
||||||
|
if self.eval_results:
|
||||||
|
return self.eval_results.pop(0)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
async def set(self, key: str, value: str, *, nx: bool, ex: int) -> bool:
|
||||||
|
self.set_calls.append((key, value, nx, ex))
|
||||||
|
return self.set_result
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeTask:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.cancelled = False
|
||||||
|
|
||||||
|
def cancel(self) -> None:
|
||||||
|
self.cancelled = True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_startup_task_coordinator_acquire_starts_refresh_for_redis_lock(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
redis = _FakeRedis(eval_results=[1])
|
||||||
|
coordinator = StartupTaskCoordinator(redis)
|
||||||
|
started: list[tuple[str, int]] = []
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
coordinator,
|
||||||
|
"_start_refresh_task",
|
||||||
|
lambda name, ttl: started.append((name, ttl)),
|
||||||
|
)
|
||||||
|
|
||||||
|
acquired = await coordinator.acquire("maintenance_scheduler", ttl=120)
|
||||||
|
|
||||||
|
assert acquired is True
|
||||||
|
assert "maintenance_scheduler" in coordinator._tokens
|
||||||
|
assert started == [("maintenance_scheduler", 120)]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_startup_task_coordinator_refresh_lock_extends_matching_token() -> None:
|
||||||
|
redis = _FakeRedis(eval_results=[1])
|
||||||
|
coordinator = StartupTaskCoordinator(redis)
|
||||||
|
coordinator._tokens["maintenance_scheduler"] = "token-1"
|
||||||
|
|
||||||
|
refreshed = await coordinator._refresh_lock("maintenance_scheduler", ttl=180)
|
||||||
|
|
||||||
|
assert refreshed is True
|
||||||
|
assert len(redis.eval_calls) == 1
|
||||||
|
_script, numkeys, key, token, ttl = redis.eval_calls[0]
|
||||||
|
assert numkeys == 1
|
||||||
|
assert key == "task_lock:maintenance_scheduler"
|
||||||
|
assert token == "token-1"
|
||||||
|
assert ttl == 180
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_startup_task_coordinator_release_cancels_refresh_task() -> None:
|
||||||
|
redis = _FakeRedis(eval_results=[1])
|
||||||
|
coordinator = StartupTaskCoordinator(redis)
|
||||||
|
coordinator._tokens["maintenance_scheduler"] = "token-1"
|
||||||
|
refresh_task = _FakeTask()
|
||||||
|
coordinator._refresh_tasks["maintenance_scheduler"] = refresh_task # type: ignore[assignment]
|
||||||
|
|
||||||
|
await coordinator.release("maintenance_scheduler")
|
||||||
|
|
||||||
|
assert refresh_task.cancelled is True
|
||||||
|
assert "maintenance_scheduler" not in coordinator._refresh_tasks
|
||||||
|
assert "maintenance_scheduler" not in coordinator._tokens
|
||||||
|
assert len(redis.eval_calls) == 1
|
||||||
|
_script, numkeys, key, token = redis.eval_calls[0]
|
||||||
|
assert numkeys == 1
|
||||||
|
assert key == "task_lock:maintenance_scheduler"
|
||||||
|
assert token == "token-1"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_startup_task_coordinator_notify_lock_lost_runs_registered_callback() -> None:
|
||||||
|
coordinator = StartupTaskCoordinator()
|
||||||
|
called: list[str] = []
|
||||||
|
|
||||||
|
async def on_lock_lost(name: str) -> None:
|
||||||
|
called.append(name)
|
||||||
|
|
||||||
|
coordinator.register_lock_lost_callback("maintenance_scheduler", on_lock_lost)
|
||||||
|
|
||||||
|
await coordinator._notify_lock_lost("maintenance_scheduler")
|
||||||
|
|
||||||
|
assert called == ["maintenance_scheduler"]
|
||||||
Reference in New Issue
Block a user