Kris
12ec5af0de
1. 新增TraceId中间件实现全链路请求追踪 2. 新增在途请求计数器与中间件支持优雅关停 3. 重构全局异常处理器，统一三模块异常响应格式 4. 优化渠道查询接口，新增多维度过滤与模糊匹配 5. 重构channels模块数据库会话管理，避免事务冲突 6. 新增目录搜索建议与目录导出功能
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
"""在途请求计数器与中间件，用于优雅关停时排空在途请求。
|
||
|
||
``InflightRequestTracker`` 通过 ``increment`` / ``decrement`` 维护在途请求
|
||
计数，计数归零时 ``wait_drained`` 立即返回；超时未归零则抛 ``TimeoutError``，
|
||
由 ``HostShutdown._waitForInflightRequests`` 捕获并记录告警后强制进入下一步
|
||
（§11.9）。
|
||
|
||
``InflightTrackingMiddleware`` 在每个 HTTP 请求进入时 ``increment``、结束时
|
||
``decrement``（``finally`` 保证异常路径也执行），确保计数准确。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
from starlette.middleware.base import BaseHTTPMiddleware
|
||
from starlette.requests import Request
|
||
|
||
__all__ = ["InflightRequestTracker", "InflightTrackingMiddleware"]
|
||
|
||
|
||
class InflightRequestTracker:
|
||
"""在途请求计数器，用于优雅关停时排空在途请求。
|
||
|
||
关键不变量：
|
||
- ``_count`` 归零时 ``_zero_event`` 被 set，``wait_drained`` 立即返回。
|
||
- ``decrement`` 防御性归零，避免计数变为负数。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._count = 0
|
||
self._zero_event = asyncio.Event()
|
||
self._zero_event.set() # 初始为 0，已触发
|
||
|
||
@property
|
||
def count(self) -> int:
|
||
"""当前在途请求数。"""
|
||
return self._count
|
||
|
||
def increment(self) -> None:
|
||
"""在途请求计数 +1，清除归零事件。"""
|
||
self._count += 1
|
||
self._zero_event.clear()
|
||
|
||
def decrement(self) -> None:
|
||
"""在途请求计数 -1，归零时触发归零事件。"""
|
||
self._count -= 1
|
||
if self._count <= 0:
|
||
self._count = 0
|
||
self._zero_event.set()
|
||
|
||
async def wait_drained(self, timeout: float) -> None:
|
||
"""等待在途请求排空。
|
||
|
||
参数：
|
||
timeout: 等待超时（秒），超时后抛 ``TimeoutError``。
|
||
|
||
Raises:
|
||
TimeoutError: 超时未排空。
|
||
"""
|
||
await asyncio.wait_for(self._zero_event.wait(), timeout=timeout)
|
||
|
||
|
||
class InflightTrackingMiddleware(BaseHTTPMiddleware):
|
||
"""在途请求计数中间件。
|
||
|
||
每个请求进入时 ``increment``、结束时 ``decrement``（``finally`` 保证
|
||
异常路径也执行）。注册时应最后注册（LIFO 顺序中最先执行），确保覆盖
|
||
所有后续中间件与路由处理。
|
||
"""
|
||
|
||
async def dispatch(self, request: Request, call_next):
|
||
tracker: InflightRequestTracker = request.app.state.inflight_tracker
|
||
tracker.increment()
|
||
try:
|
||
return await call_next(request)
|
||
finally:
|
||
tracker.decrement()
|