feat: 添加启动任务开关,修复统计聚合内存泄漏

- 新增 CACHE_WARMUP_ENABLED 和 MAINTENANCE_STARTUP_TASKS_ENABLED 环境变量,
  允许禁用缓存预热和维护调度器启动任务
- 在统计聚合批量处理循环中添加 db.expunge_all(),释放 Session identity map,
  防止 ORM 对象累积导致内存暴涨
- 添加启动任务开关的单元测试
This commit is contained in:
fawney19
2026-03-10 13:42:38 +08:00
parent 2b21a75982
commit f40e8037dd
5 changed files with 111 additions and 1 deletions

View File

@@ -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
# - 未设置: 开发环境启用,生产环境禁用

View File

@@ -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())

View File

@@ -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}")

View File

@@ -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