71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
import threading
|
|
import time
|
|
from functools import wraps
|
|
from flask import current_app, jsonify, stream_with_context, Response
|
|
|
|
class ExecutionTimeoutMonitor:
|
|
"""监控请求执行时间，超时后释放信号量"""
|
|
def __init__(self, timeout, semaphore):
|
|
self.timeout = timeout
|
|
self.semaphore = semaphore
|
|
self.is_timeout = False
|
|
self.thread = threading.Thread(target=self._monitor)
|
|
self.thread.daemon = True
|
|
|
|
def _monitor(self):
|
|
"""等待指定时间后标记为超时并释放信号量
|
|
目前的超时机制是通过 ExecutionTimeoutMonitor 的标志位来通知生成器超时。
|
|
但是，生成器只有在下一次 yield 时才会检查标志位，因此不会在当前执行的代码段中间强行中断。
|
|
"""
|
|
time.sleep(self.timeout)
|
|
self.is_timeout = True
|
|
self.semaphore.release() # 超时后释放信号量
|
|
current_app.logger.error(f"Request execution exceeded {self.timeout} seconds and was terminated.")
|
|
|
|
def start(self):
|
|
self.thread.start()
|
|
|
|
def require_connection_limit(timeout=900):
|
|
"""装饰器：确保路由使用连接限制，并监控请求执行时间"""
|
|
def decorator(f):
|
|
@wraps(f)
|
|
def wrapped(*args, **kwargs):
|
|
limiter = getattr(current_app, 'connection_limiter', None)
|
|
if limiter is None:
|
|
current_app.logger.error("ConnectionLimiter 未初始化")
|
|
return jsonify({'error': 'Server configuration error'}), 500
|
|
|
|
# 阻塞方式获取信号量
|
|
limiter.semaphore.acquire()
|
|
try:
|
|
# 启动执行超时监控器
|
|
monitor = ExecutionTimeoutMonitor(timeout, limiter.semaphore)
|
|
monitor.start()
|
|
|
|
generator = f(*args, **kwargs)
|
|
|
|
@stream_with_context
|
|
def generator_wrapper():
|
|
try:
|
|
for item in generator:
|
|
if monitor.is_timeout:
|
|
# 如果已超时，终止生成器执行
|
|
current_app.logger.error("Request exceeded execution time and was terminated.")
|
|
break
|
|
yield item
|
|
finally:
|
|
if not monitor.is_timeout:
|
|
limiter.semaphore.release() # 正常结束时释放信号量
|
|
|
|
return Response(generator_wrapper(), mimetype='text/event-stream')
|
|
except Exception as e:
|
|
limiter.semaphore.release() # 异常时释放信号量
|
|
current_app.logger.error(f"Exception in route: {e}")
|
|
return jsonify({'error': 'Internal server error'}), 500
|
|
return wrapped
|
|
return decorator
|
|
|
|
class ConnectionLimiter:
|
|
def __init__(self, max_connections=1):
|
|
self.semaphore = threading.Semaphore(max_connections)
|