357 lines
11 KiB
Python
357 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
重试机制模块
|
||
提供重试装饰器和工具函数
|
||
"""
|
||
import time
|
||
import logging
|
||
from functools import wraps
|
||
from typing import Callable, Optional, Type, Tuple, Any
|
||
|
||
from src.logging_config import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
def retry(
|
||
max_attempts: int = 3,
|
||
delay: float = 1.0,
|
||
backoff_factor: float = 2.0,
|
||
exceptions: Optional[Tuple[Type[Exception], ...]] = None,
|
||
on_retry: Optional[Callable[[int, Exception], None]] = None
|
||
) -> Callable:
|
||
"""
|
||
重试装饰器
|
||
|
||
参数:
|
||
max_attempts: 最大重试次数
|
||
delay: 初始延迟时间（秒）
|
||
backoff_factor: 退避因子，每次重试延迟时间乘以该因子
|
||
exceptions: 要捕获的异常类型，None表示捕获所有异常
|
||
on_retry: 重试时的回调函数，参数为 (attempt, exception)
|
||
|
||
使用示例:
|
||
@retry(max_attempts=3, delay=2.0, backoff_factor=2.0)
|
||
def fetch_data():
|
||
# 可能失败的代码
|
||
pass
|
||
|
||
@retry(max_attempts=5, exceptions=(ConnectionError, TimeoutError))
|
||
def network_request():
|
||
# 网络请求代码
|
||
pass
|
||
"""
|
||
def decorator(func: Callable) -> Callable:
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs) -> Any:
|
||
last_exception = None
|
||
|
||
for attempt in range(max_attempts):
|
||
try:
|
||
return func(*args, **kwargs)
|
||
except Exception as e:
|
||
# 检查是否需要捕获此异常
|
||
if exceptions and not isinstance(e, exceptions):
|
||
raise
|
||
|
||
last_exception = e
|
||
|
||
# 如果是最后一次尝试，不再重试
|
||
if attempt == max_attempts - 1:
|
||
logger.error(
|
||
f"{func.__name__} 在 {max_attempts} 次尝试后仍然失败: {e}"
|
||
)
|
||
raise
|
||
|
||
# 计算延迟时间
|
||
current_delay = delay * (backoff_factor ** attempt)
|
||
|
||
logger.warning(
|
||
f"{func.__name__} 第 {attempt + 1} 次尝试失败: {e}, "
|
||
f"{current_delay:.2f}秒后重试..."
|
||
)
|
||
|
||
# 调用重试回调
|
||
if on_retry:
|
||
try:
|
||
on_retry(attempt + 1, e)
|
||
except Exception as callback_error:
|
||
logger.error(f"重试回调执行失败: {callback_error}")
|
||
|
||
# 等待
|
||
time.sleep(current_delay)
|
||
|
||
# 理论上不会到达这里，但为了类型检查
|
||
if last_exception:
|
||
raise last_exception
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
def retry_with_exponential_backoff(
|
||
max_attempts: int = 3,
|
||
initial_delay: float = 1.0,
|
||
max_delay: float = 60.0
|
||
) -> Callable:
|
||
"""
|
||
使用指数退避的重试装饰器
|
||
|
||
参数:
|
||
max_attempts: 最大重试次数
|
||
initial_delay: 初始延迟时间（秒）
|
||
max_delay: 最大延迟时间（秒）
|
||
|
||
使用示例:
|
||
@retry_with_exponential_backoff(max_attempts=5, initial_delay=2.0)
|
||
def api_call():
|
||
# API调用代码
|
||
pass
|
||
"""
|
||
def decorator(func: Callable) -> Callable:
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs) -> Any:
|
||
last_exception = None
|
||
|
||
for attempt in range(max_attempts):
|
||
try:
|
||
return func(*args, **kwargs)
|
||
except Exception as e:
|
||
last_exception = e
|
||
|
||
if attempt == max_attempts - 1:
|
||
logger.error(
|
||
f"{func.__name__} 在 {max_attempts} 次尝试后仍然失败: {e}"
|
||
)
|
||
raise
|
||
|
||
# 计算延迟时间（指数退避，但不超过最大延迟）
|
||
current_delay = min(initial_delay * (2 ** attempt), max_delay)
|
||
|
||
logger.warning(
|
||
f"{func.__name__} 第 {attempt + 1} 次尝试失败: {e}, "
|
||
f"{current_delay:.2f}秒后重试..."
|
||
)
|
||
|
||
time.sleep(current_delay)
|
||
|
||
if last_exception:
|
||
raise last_exception
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
def retry_on_exception(
|
||
exception_type: Type[Exception],
|
||
max_attempts: int = 3,
|
||
delay: float = 1.0
|
||
) -> Callable:
|
||
"""
|
||
只在特定异常时重试的装饰器
|
||
|
||
参数:
|
||
exception_type: 要捕获的异常类型
|
||
max_attempts: 最大重试次数
|
||
delay: 延迟时间（秒）
|
||
|
||
使用示例:
|
||
@retry_on_exception(ConnectionError, max_attempts=5, delay=2.0)
|
||
def fetch_data():
|
||
# 可能抛出 ConnectionError 的代码
|
||
pass
|
||
"""
|
||
def decorator(func: Callable) -> Callable:
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs) -> Any:
|
||
last_exception = None
|
||
|
||
for attempt in range(max_attempts):
|
||
try:
|
||
return func(*args, **kwargs)
|
||
except exception_type as e:
|
||
last_exception = e
|
||
|
||
if attempt == max_attempts - 1:
|
||
logger.error(
|
||
f"{func.__name__} 在 {max_attempts} 次尝试后仍然失败: {e}"
|
||
)
|
||
raise
|
||
|
||
logger.warning(
|
||
f"{func.__name__} 第 {attempt + 1} 次尝试失败: {e}, "
|
||
f"{delay:.2f}秒后重试..."
|
||
)
|
||
|
||
time.sleep(delay)
|
||
|
||
if last_exception:
|
||
raise last_exception
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
class RetryContext:
|
||
"""重试上下文管理器"""
|
||
|
||
def __init__(
|
||
self,
|
||
operation_name: str,
|
||
max_attempts: int = 3,
|
||
delay: float = 1.0,
|
||
backoff_factor: float = 2.0,
|
||
exceptions: Optional[Tuple[Type[Exception], ...]] = None
|
||
):
|
||
"""
|
||
初始化重试上下文
|
||
|
||
参数:
|
||
operation_name: 操作名称
|
||
max_attempts: 最大重试次数
|
||
delay: 初始延迟时间（秒）
|
||
backoff_factor: 退避因子
|
||
exceptions: 要捕获的异常类型
|
||
"""
|
||
self.operation_name = operation_name
|
||
self.max_attempts = max_attempts
|
||
self.delay = delay
|
||
self.backoff_factor = backoff_factor
|
||
self.exceptions = exceptions
|
||
self.attempt = 0
|
||
|
||
def __enter__(self):
|
||
self.attempt = 0
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
if exc_type is None:
|
||
return False
|
||
|
||
# 检查是否需要捕获此异常
|
||
if self.exceptions and not isinstance(exc_val, self.exceptions):
|
||
return False
|
||
|
||
self.attempt += 1
|
||
|
||
# 如果超过最大尝试次数，不再重试
|
||
if self.attempt >= self.max_attempts:
|
||
logger.error(
|
||
f"{self.operation_name} 在 {self.max_attempts} 次尝试后仍然失败: {exc_val}"
|
||
)
|
||
return False
|
||
|
||
# 计算延迟时间
|
||
current_delay = self.delay * (self.backoff_factor ** (self.attempt - 1))
|
||
|
||
logger.warning(
|
||
f"{self.operation_name} 第 {self.attempt} 次尝试失败: {exc_val}, "
|
||
f"{current_delay:.2f}秒后重试..."
|
||
)
|
||
|
||
# 等待
|
||
time.sleep(current_delay)
|
||
|
||
# 抑制异常，继续重试
|
||
return True
|
||
|
||
|
||
def async_retry(
|
||
max_attempts: int = 3,
|
||
delay: float = 1.0,
|
||
backoff_factor: float = 2.0
|
||
) -> Callable:
|
||
"""
|
||
异步重试装饰器（用于异步函数）
|
||
|
||
参数:
|
||
max_attempts: 最大重试次数
|
||
delay: 初始延迟时间（秒）
|
||
backoff_factor: 退避因子
|
||
|
||
使用示例:
|
||
@async_retry(max_attempts=3, delay=2.0)
|
||
async def async_fetch_data():
|
||
# 异步代码
|
||
pass
|
||
"""
|
||
import asyncio
|
||
|
||
def decorator(func: Callable) -> Callable:
|
||
@wraps(func)
|
||
async def wrapper(*args, **kwargs) -> Any:
|
||
last_exception = None
|
||
|
||
for attempt in range(max_attempts):
|
||
try:
|
||
return await func(*args, **kwargs)
|
||
except Exception as e:
|
||
last_exception = e
|
||
|
||
if attempt == max_attempts - 1:
|
||
logger.error(
|
||
f"{func.__name__} 在 {max_attempts} 次尝试后仍然失败: {e}"
|
||
)
|
||
raise
|
||
|
||
# 计算延迟时间
|
||
current_delay = delay * (backoff_factor ** attempt)
|
||
|
||
logger.warning(
|
||
f"{func.__name__} 第 {attempt + 1} 次尝试失败: {e}, "
|
||
f"{current_delay:.2f}秒后重试..."
|
||
)
|
||
|
||
# 异步等待
|
||
await asyncio.sleep(current_delay)
|
||
|
||
if last_exception:
|
||
raise last_exception
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 测试代码
|
||
|
||
# 测试重试装饰器
|
||
call_count = 0
|
||
|
||
@retry(max_attempts=3, delay=0.1)
|
||
def test_retry():
|
||
global call_count
|
||
call_count += 1
|
||
print(f"调用次数: {call_count}")
|
||
if call_count < 3:
|
||
raise ValueError("测试异常")
|
||
return "成功"
|
||
|
||
result = test_retry()
|
||
print(f"测试结果: {result}")
|
||
|
||
# 测试重试上下文管理器
|
||
context_call_count = 0
|
||
|
||
def test_context_operation():
|
||
global context_call_count
|
||
context_call_count += 1
|
||
print(f"上下文调用次数: {context_call_count}")
|
||
if context_call_count < 3:
|
||
raise ValueError("测试异常")
|
||
return "成功"
|
||
|
||
with RetryContext("测试操作", max_attempts=3, delay=0.1):
|
||
result = test_context_operation()
|
||
print(f"上下文测试结果: {result}")
|
||
|
||
# 测试特定异常重试
|
||
@retry_on_exception(ValueError, max_attempts=3, delay=0.1)
|
||
def test_specific_exception():
|
||
raise ValueError("测试异常")
|
||
|
||
try:
|
||
test_specific_exception()
|
||
except ValueError as e:
|
||
print(f"特定异常测试: {e}")
|