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