error_handling.py

"""
Error Handling and Retry Mechanisms
错误处理和重试机制，包括断路器、超时处理、死信队列等
"""

import asyncio
import logging
import time
from enum import Enum
from typing import Dict, Any, Optional, Callable, List
from dataclasses import dataclass, field
from datetime import datetime, timedelta
import json

from message_queue_client import Message

logger = logging.getLogger(__name__)


class CircuitState(Enum):
    """断路器状态"""
    CLOSED = "closed"      # 正常状态
    OPEN = "open"          # 断路状态
    HALF_OPEN = "half_open" # 半开状态


@dataclass
class CircuitBreakerConfig:
    """断路器配置"""
    failure_threshold: int = 5          # 失败阈值
    recovery_timeout: int = 60          # 恢复超时时间（秒）
    success_threshold: int = 3          # 成功阈值（半开状态）
    failure_count: int = 0              # 当前失败计数
    success_count: int = 0              # 当前成功计数
    last_failure_time: Optional[float] = None  # 上次失败时间
    state: CircuitState = CircuitState.CLOSED   # 当前状态


class CircuitBreaker:
    """断路器模式实现"""

    def __init__(self, config: CircuitBreakerConfig):
        self.config = config
        self._lock = asyncio.Lock()

    async def can_execute(self) -> bool:
        """检查是否可以执行请求"""
        async with self._lock:
            current_time = time.time()

            if self.config.state == CircuitState.CLOSED:
                return True

            elif self.config.state == CircuitState.OPEN:
                # 检查是否到了恢复时间
                if (self.config.last_failure_time and
                    current_time - self.config.last_failure_time >= self.config.recovery_timeout):
                    self.config.state = CircuitState.HALF_OPEN
                    self.config.success_count = 0
                    logger.info("Circuit breaker moved to HALF_OPEN state")
                    return True
                return False

            elif self.config.state == CircuitState.HALF_OPEN:
                return True

            return False

    async def record_success(self):
        """记录成功"""
        async with self._lock:
            if self.config.state == CircuitState.HALF_OPEN:
                self.config.success_count += 1
                if self.config.success_count >= self.config.success_threshold:
                    self.config.state = CircuitState.CLOSED
                    self.config.failure_count = 0
                    logger.info("Circuit breaker moved to CLOSED state")
            elif self.config.state == CircuitState.CLOSED:
                self.config.failure_count = 0

    async def record_failure(self):
        """记录失败"""
        async with self._lock:
            self.config.failure_count += 1
            self.config.last_failure_time = time.time()

            if self.config.state == CircuitState.HALF_OPEN:
                self.config.state = CircuitState.OPEN
                logger.warning("Circuit breaker moved to OPEN state (half-open failure)")
            elif (self.config.state == CircuitState.CLOSED and
                  self.config.failure_count >= self.config.failure_threshold):
                self.config.state = CircuitState.OPEN
                logger.warning("Circuit breaker moved to OPEN state (threshold exceeded)")


@dataclass
class RetryConfig:
    """重试配置"""
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True


class RetryHandler:
    """重试处理器"""

    def __init__(self, config: RetryConfig):
        self.config = config

    def calculate_delay(self, attempt: int) -> float:
        """计算重试延迟"""
        delay = self.config.base_delay * (self.config.exponential_base ** attempt)
        delay = min(delay, self.config.max_delay)

        if self.config.jitter:
            # 添加随机抖动
            import random
            delay = delay * (0.5 + random.random() * 0.5)

        return delay

    async def execute_with_retry(self, func: Callable, *args, **kwargs) -> Any:
        """执行带重试的函数"""
        last_exception = None

        for attempt in range(self.config.max_attempts):
            try:
                result = await func(*args, **kwargs)
                return result

            except Exception as e:
                last_exception = e
                logger.warning(f"Attempt {attempt + 1} failed: {e}")

                if attempt < self.config.max_attempts - 1:
                    delay = self.calculate_delay(attempt)
                    logger.info(f"Retrying in {delay:.2f} seconds...")
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"All retry attempts failed after {self.config.max_attempts} tries")
                    raise last_exception

        raise last_exception


class DeadLetterQueue:
    """死信队列处理"""

    def __init__(self, queue_client, max_retries: int = 3):
        self.queue_client = queue_client
        self.max_retries = max_retries

    def should_send_to_dlq(self, message: Message) -> bool:
        """判断是否应该发送到死信队列"""
        retry_count = message.metadata.get("retry_count", 0)
        return retry_count >= self.max_retries

    async def handle_failed_message(self, message: Message, error: str):
        """处理失败的消息"""
        try:
            # 增加重试计数
            message.metadata["retry_count"] = message.metadata.get("retry_count", 0) + 1

            if self.should_send_to_dlq(message):
                # 发送到死信队列
                dlq_message = Message(
                    agent_name=f"{message.agent_name}.dlq",
                    session_id=message.session_id,
                    payload={
                        "original_message": message.to_dict(),
                        "error": error,
                        "failed_at": datetime.utcnow().isoformat()
                    }
                )
                self.queue_client.publish(f"{message.agent_name}.dlq", dlq_message)
                logger.warning(f"Message {message.message_id} sent to DLQ")
            else:
                # 重新排队等待重试
                logger.info(f"Message {message.message_id} requeued for retry (attempt {message.metadata['retry_count']})")
                self.queue_client.publish(f"{message.agent_name}.requests", message)

        except Exception as e:
            logger.error(f"Failed to handle failed message: {e}")


class ErrorHandler:
    """统一错误处理器"""

    def __init__(self, circuit_breaker: CircuitBreaker, retry_handler: RetryHandler,
                 dead_letter_queue: DeadLetterQueue):
        self.circuit_breaker = circuit_breaker
        self.retry_handler = retry_handler
        self.dead_letter_queue = dead_letter_queue

    async def execute_with_protection(self, func: Callable, message: Message, *args, **kwargs) -> Any:
        """在保护下执行函数"""
        try:
            # 检查断路器状态
            if not await self.circuit_breaker.can_execute():
                raise Exception("Circuit breaker is OPEN")

            # 使用重试机制执行
            result = await self.retry_handler.execute_with_retry(func, *args, **kwargs)

            # 记录成功
            await self.circuit_breaker.record_success()
            return result

        except Exception as e:
            # 记录失败
            await self.circuit_breaker.record_failure()

            # 处理失败的消息
            await self.dead_letter_queue.handle_failed_message(message, str(e))

            logger.error(f"Execution failed for message {message.message_id}: {e}")
            raise


class HealthMonitor:
    """健康监控器"""

    def __init__(self, service_name: str, check_interval: int = 30):
        self.service_name = service_name
        self.check_interval = check_interval
        self.health_status = {
            "status": "unknown",
            "last_check": None,
            "consecutive_failures": 0,
            "total_checks": 0,
            "successful_checks": 0
        }
        self._running = False

    async def start_monitoring(self, health_check_func: Callable):
        """开始健康监控"""
        self._running = True
        while self._running:
            try:
                await self._perform_health_check(health_check_func)
                await asyncio.sleep(self.check_interval)
            except Exception as e:
                logger.error(f"Health monitoring error: {e}")
                await asyncio.sleep(self.check_interval)

    async def _perform_health_check(self, health_check_func: Callable):
        """执行健康检查"""
        try:
            self.health_status["total_checks"] += 1
            is_healthy = await health_check_func()

            if is_healthy:
                self.health_status["status"] = "healthy"
                self.health_status["successful_checks"] += 1
                self.health_status["consecutive_failures"] = 0
            else:
                self.health_status["status"] = "unhealthy"
                self.health_status["consecutive_failures"] += 1

        except Exception as e:
            self.health_status["status"] = "error"
            self.health_status["consecutive_failures"] += 1
            logger.error(f"Health check failed: {e}")

        finally:
            self.health_status["last_check"] = datetime.utcnow().isoformat()

    def get_health_status(self) -> Dict[str, Any]:
        """获取健康状态"""
        return self.health_status.copy()

    def stop_monitoring(self):
        """停止健康监控"""
        self._running = False


class MetricsCollector:
    """指标收集器"""

    def __init__(self, service_name: str):
        self.service_name = service_name
        self.metrics = {
            "requests_total": 0,
            "requests_success": 0,
            "requests_failed": 0,
            "request_duration_seconds": [],
            "errors_by_type": {},
            "start_time": datetime.utcnow().isoformat()
        }

    def record_request(self, success: bool, duration: float, error_type: Optional[str] = None):
        """记录请求指标"""
        self.metrics["requests_total"] += 1
        self.metrics["request_duration_seconds"].append(duration)

        if success:
            self.metrics["requests_success"] += 1
        else:
            self.metrics["requests_failed"] += 1
            if error_type:
                self.metrics["errors_by_type"][error_type] = self.metrics["errors_by_type"].get(error_type, 0) + 1

    def get_metrics(self) -> Dict[str, Any]:
        """获取指标数据"""
        metrics = self.metrics.copy()
        if metrics["request_duration_seconds"]:
            durations = metrics["request_duration_seconds"]
            metrics["request_duration_avg"] = sum(durations) / len(durations)
            metrics["request_duration_p95"] = sorted(durations)[int(len(durations) * 0.95)]
        return metrics


# 集成错误处理的消息队列节点
class ResilientMessageQueueNode:
    """具有弹性的消息队列节点"""

    def __init__(self, node_name: str, queue_client,
                 circuit_breaker_config: CircuitBreakerConfig = None,
                 retry_config: RetryConfig = None,
                 timeout: int = 30):
        self.node_name = node_name
        self.queue_client = queue_client
        self.timeout = timeout

        # 初始化错误处理组件
        self.circuit_breaker = CircuitBreaker(circuit_breaker_config or CircuitBreakerConfig())
        self.retry_handler = RetryHandler(retry_config or RetryConfig())
        self.dead_letter_queue = DeadLetterQueue(queue_client)
        self.error_handler = ErrorHandler(self.circuit_breaker, self.retry_handler, self.dead_letter_queue)
        self.health_monitor = HealthMonitor(node_name)
        self.metrics_collector = MetricsCollector(node_name)

    async def resilient_process(self, message: Message, process_func: Callable) -> Any:
        """弹性处理消息"""
        start_time = time.time()

        try:
            # 在保护下执行处理函数
            result = await self.error_handler.execute_with_protection(
                process_func, message, message.payload
            )

            # 记录成功指标
            duration = time.time() - start_time
            self.metrics_collector.record_request(True, duration)

            return result

        except Exception as e:
            # 记录失败指标
            duration = time.time() - start_time
            self.metrics_collector.record_request(False, duration, type(e).__name__)

            logger.error(f"Resilient processing failed for message {message.message_id}: {e}")
            raise

    def get_status(self) -> Dict[str, Any]:
        """获取节点状态"""
        return {
            "node_name": self.node_name,
            "circuit_breaker_state": self.circuit_breaker.config.state.value,
            "health_status": self.health_monitor.get_health_status(),
            "metrics": self.metrics_collector.get_metrics()
        }


# 使用示例
async def example_error_handling():
    """错误处理示例"""
    from message_queue_client import create_message_queue_client

    # 创建消息队列客户端
    queue_client = create_message_queue_client("redis", host="localhost", port=6379)

    # 创建弹性节点
    node = ResilientMessageQueueNode(
        "test_node",
        queue_client,
        CircuitBreakerConfig(failure_threshold=3, recovery_timeout=30),
        RetryConfig(max_attempts=3, base_delay=1.0)
    )

    # 模拟消息
    test_message = Message(
        agent_name="test_agent",
        session_id="test_session",
        payload={"test": "data"}
    )

    # 模拟处理函数
    async def mock_process(data):
        # 随机失败用于测试
        import random
        if random.random() < 0.5:
            raise Exception("Mock processing error")
        return {"result": "success"}

    try:
        # 在弹性保护下处理
        result = await node.resilient_process(test_message, mock_process)
        print(f"Processing result: {result}")
    except Exception as e:
        print(f"Processing failed after all retries: {e}")

    # 获取状态
    status = node.get_status()
    print(f"Node status: {json.dumps(status, indent=2)}")

    queue_client.close()


if __name__ == "__main__":
    asyncio.run(example_error_handling())