65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
|
from fastapi import Request, Response
|
||
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
||
|
|
from starlette.types import ASGIApp
|
||
|
|
import time
|
||
|
|
from typing import Callable
|
||
|
|
from ..utils import get_logger
|
||
|
|
|
||
|
|
logger = get_logger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class PerformanceMiddleware(BaseHTTPMiddleware):
|
||
|
|
"""性能监控中间件"""
|
||
|
|
|
||
|
|
# 类变量，用于存储实例引用
|
||
|
|
_instance = None
|
||
|
|
|
||
|
|
def __init__(self, app: ASGIApp):
|
||
|
|
super().__init__(app)
|
||
|
|
self.logger = get_logger("performance")
|
||
|
|
self.total_requests = 0
|
||
|
|
self.requests_per_status = {}
|
||
|
|
self.total_response_time = 0
|
||
|
|
|
||
|
|
# 保存实例引用
|
||
|
|
PerformanceMiddleware._instance = self
|
||
|
|
|
||
|
|
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
||
|
|
"""处理请求并添加性能头"""
|
||
|
|
start_time = time.time()
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = await call_next(request)
|
||
|
|
response_time = time.time() - start_time
|
||
|
|
|
||
|
|
# 添加性能头
|
||
|
|
response.headers["X-Response-Time"] = f"{response_time:.4f}"
|
||
|
|
|
||
|
|
# 更新请求计数
|
||
|
|
self.total_requests += 1
|
||
|
|
status_code = response.status_code
|
||
|
|
if status_code not in self.requests_per_status:
|
||
|
|
self.requests_per_status[status_code] = 0
|
||
|
|
self.requests_per_status[status_code] += 1
|
||
|
|
|
||
|
|
# 更新总响应时间
|
||
|
|
self.total_response_time += response_time
|
||
|
|
|
||
|
|
return response
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
raise
|
||
|
|
|
||
|
|
def get_stats(self):
|
||
|
|
"""获取性能统计信息"""
|
||
|
|
return {
|
||
|
|
"total_requests": self.total_requests,
|
||
|
|
"requests_per_status": self.requests_per_status,
|
||
|
|
"total_response_time": self.total_response_time,
|
||
|
|
"average_response_time": self.total_response_time / self.total_requests if self.total_requests > 0 else 0
|
||
|
|
}
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def get_instance(cls):
|
||
|
|
"""获取中间件实例"""
|
||
|
|
return cls._instance
|