n1

2026-02-01 20:56:37 +08:00
parent 486313044c
commit 0a576b04cf
28 changed files with 12410 additions and 6701 deletions
--- a/src/retry.py
+++ b/src/retry.py
@@ -0,0 +1,356 @@
+#!/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}")