mirror of
https://github.com/fawney19/Aether.git
synced 2026-01-03 00:02:28 +08:00
b579420690
- Improve database pool capacity logging with detailed configuration parameters - Optimize database session dependency injection with middleware-managed lifecycle - Simplify plugin middleware by delegating session creation to FastAPI dependencies - Fix import path in auth routes (relative to absolute) - Add safety checks for database session management across middleware exception handlers - Ensure session cleanup only when not managed by middleware (avoid premature cleanup)
121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
"""
|
||
批量提交器 - 减少数据库 commit 次数，提升并发能力
|
||
|
||
核心思想：
|
||
- 非关键数据（监控、统计）不立即 commit
|
||
- 在后台定期批量 commit
|
||
- 关键数据（计费）仍然立即 commit
|
||
"""
|
||
|
||
import asyncio
|
||
from typing import Set
|
||
|
||
from src.core.logger import logger
|
||
from sqlalchemy.orm import Session
|
||
|
||
|
||
class BatchCommitter:
|
||
"""批量提交管理器"""
|
||
|
||
def __init__(self, interval_seconds: float = 1.0):
|
||
"""
|
||
Args:
|
||
interval_seconds: 批量提交间隔（秒）
|
||
"""
|
||
self.interval_seconds = interval_seconds
|
||
self._pending_sessions: Set[Session] = set()
|
||
self._lock = asyncio.Lock()
|
||
self._task = None
|
||
|
||
async def start(self):
|
||
"""启动后台批量提交任务"""
|
||
if self._task is None:
|
||
self._task = asyncio.create_task(self._batch_commit_loop())
|
||
logger.info(f"批量提交器已启动，间隔: {self.interval_seconds}s")
|
||
|
||
async def stop(self):
|
||
"""停止后台任务"""
|
||
if self._task:
|
||
self._task.cancel()
|
||
try:
|
||
await self._task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._task = None
|
||
logger.info("批量提交器已停止")
|
||
|
||
def mark_dirty(self, session: Session):
|
||
"""标记 Session 有待提交的更改"""
|
||
# 请求级事务由中间件统一 commit/rollback；避免后台任务在请求中途误提交。
|
||
if session is None:
|
||
return
|
||
if session.info.get("managed_by_middleware"):
|
||
return
|
||
self._pending_sessions.add(session)
|
||
|
||
async def _batch_commit_loop(self):
|
||
"""后台批量提交循环"""
|
||
while True:
|
||
try:
|
||
await asyncio.sleep(self.interval_seconds)
|
||
await self._commit_all()
|
||
except asyncio.CancelledError:
|
||
# 关闭前提交所有待处理的
|
||
await self._commit_all()
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"批量提交出错: {e}")
|
||
|
||
async def _commit_all(self):
|
||
"""提交所有待处理的 Session"""
|
||
async with self._lock:
|
||
if not self._pending_sessions:
|
||
return
|
||
|
||
sessions_to_commit = list(self._pending_sessions)
|
||
self._pending_sessions.clear()
|
||
|
||
committed = 0
|
||
failed = 0
|
||
|
||
for session in sessions_to_commit:
|
||
try:
|
||
session.commit()
|
||
committed += 1
|
||
except Exception as e:
|
||
logger.error(f"提交 Session 失败: {e}")
|
||
try:
|
||
session.rollback()
|
||
except:
|
||
pass
|
||
failed += 1
|
||
|
||
if committed > 0:
|
||
logger.debug(f"批量提交完成: {committed} 个 Session")
|
||
if failed > 0:
|
||
logger.warning(f"批量提交失败: {failed} 个 Session")
|
||
|
||
|
||
# 全局单例
|
||
_batch_committer: BatchCommitter = None
|
||
|
||
|
||
def get_batch_committer() -> BatchCommitter:
|
||
"""获取全局批量提交器"""
|
||
global _batch_committer
|
||
if _batch_committer is None:
|
||
_batch_committer = BatchCommitter(interval_seconds=1.0)
|
||
return _batch_committer
|
||
|
||
|
||
async def init_batch_committer():
|
||
"""初始化并启动批量提交器"""
|
||
committer = get_batch_committer()
|
||
await committer.start()
|
||
|
||
|
||
async def shutdown_batch_committer():
|
||
"""关闭批量提交器"""
|
||
committer = get_batch_committer()
|
||
await committer.stop()
|