mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat: 添加启动任务开关,修复统计聚合内存泄漏
- 新增 CACHE_WARMUP_ENABLED 和 MAINTENANCE_STARTUP_TASKS_ENABLED 环境变量, 允许禁用缓存预热和维护调度器启动任务 - 在统计聚合批量处理循环中添加 db.expunge_all(),释放 Session identity map, 防止 ORM 对象累积导致内存暴涨 - 添加启动任务开关的单元测试
This commit is contained in:
@@ -312,6 +312,14 @@ class Config:
|
||||
# 每个用户最多可创建的 Management Token 数量
|
||||
self.management_token_max_per_user = int(os.getenv("MANAGEMENT_TOKEN_MAX_PER_USER", "20"))
|
||||
|
||||
# 启动任务开关
|
||||
# CACHE_WARMUP_ENABLED: 是否在启动时执行缓存预热
|
||||
# MAINTENANCE_STARTUP_TASKS_ENABLED: 是否在启动时执行维护调度器初始化任务(清理、统计回填等)
|
||||
self.cache_warmup_enabled = os.getenv("CACHE_WARMUP_ENABLED", "true").lower() == "true"
|
||||
self.maintenance_startup_tasks_enabled = (
|
||||
os.getenv("MAINTENANCE_STARTUP_TASKS_ENABLED", "true").lower() == "true"
|
||||
)
|
||||
|
||||
# API 文档配置
|
||||
# DOCS_ENABLED: 是否启用 API 文档(/docs, /redoc, /openapi.json)
|
||||
# - 未设置: 开发环境启用,生产环境禁用
|
||||
|
||||
@@ -17,6 +17,7 @@ from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.enums import UserRole
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
@@ -172,4 +173,8 @@ class CacheWarmupService:
|
||||
|
||||
async def start_cache_warmup() -> None:
|
||||
"""启动缓存预热(作为后台任务)"""
|
||||
if not config.cache_warmup_enabled:
|
||||
logger.info("缓存预热已禁用(CACHE_WARMUP_ENABLED=false)")
|
||||
return
|
||||
|
||||
asyncio.create_task(CacheWarmupService.warmup_all())
|
||||
|
||||
@@ -23,6 +23,7 @@ from typing import Any
|
||||
|
||||
from sqlalchemy import delete, literal_column, text
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import AuditLog, Provider, RequestCandidate, Usage
|
||||
@@ -232,7 +233,10 @@ class MaintenanceScheduler:
|
||||
)
|
||||
|
||||
# 启动时执行一次初始化任务
|
||||
asyncio.create_task(self._run_startup_tasks())
|
||||
if config.maintenance_startup_tasks_enabled:
|
||||
asyncio.create_task(self._run_startup_tasks())
|
||||
else:
|
||||
logger.info("维护调度器启动任务已禁用(MAINTENANCE_STARTUP_TASKS_ENABLED=false)")
|
||||
|
||||
async def _run_startup_tasks(self) -> None:
|
||||
"""启动时执行的初始化任务"""
|
||||
@@ -470,6 +474,7 @@ class MaintenanceScheduler:
|
||||
StatsAggregatorService.aggregate_daily_stats_bundle(
|
||||
db, current_date_utc, user_ids=user_ids
|
||||
)
|
||||
db.expunge_all()
|
||||
except Exception as e:
|
||||
failed_dates += 1
|
||||
logger.warning(f"回填日期 {current_date} 失败: {e}")
|
||||
|
||||
@@ -1096,6 +1096,7 @@ class StatsAggregatorService:
|
||||
current_date = start_date
|
||||
while current_date < today_utc:
|
||||
StatsAggregatorService.aggregate_daily_stats_bundle(db, current_date, user_ids=user_ids)
|
||||
db.expunge_all() # 释放 Session identity map,防止 ORM 对象累积导致内存暴涨
|
||||
count += 1
|
||||
current_date += timedelta(days=1)
|
||||
|
||||
@@ -1114,6 +1115,7 @@ class StatsAggregatorService:
|
||||
count = 0
|
||||
while current <= end_dt:
|
||||
StatsAggregatorService.aggregate_daily_api_key_stats(db, current, commit=True)
|
||||
db.expunge_all()
|
||||
count += 1
|
||||
current += timedelta(days=1)
|
||||
return count
|
||||
@@ -1145,9 +1147,11 @@ class StatsAggregatorService:
|
||||
processed += 1
|
||||
if processed % 30 == 0:
|
||||
db.commit()
|
||||
db.expunge_all()
|
||||
current += timedelta(days=1)
|
||||
|
||||
db.commit()
|
||||
db.expunge_all()
|
||||
return processed
|
||||
|
||||
@staticmethod
|
||||
@@ -1183,6 +1187,7 @@ class StatsAggregatorService:
|
||||
db.commit()
|
||||
last_id = records[-1].id
|
||||
total_processed += len(records)
|
||||
db.expunge_all()
|
||||
|
||||
return total_processed
|
||||
|
||||
@@ -1194,6 +1199,7 @@ class StatsAggregatorService:
|
||||
count = 0
|
||||
while current <= end_dt:
|
||||
StatsAggregatorService.aggregate_daily_error_stats(db, current, commit=True)
|
||||
db.expunge_all()
|
||||
count += 1
|
||||
current += timedelta(days=1)
|
||||
return count
|
||||
|
||||
Reference in New Issue
Block a user