gogo2/COBY/connectors/circuit_breaker.py

"""
Circuit breaker pattern implementation for exchange connections.
"""

import time
from enum import Enum
from typing import Optional, Callable, Any
from ..utils.logging import get_logger

logger = get_logger(__name__)


class CircuitState(Enum):
    """Circuit breaker states"""
    CLOSED = "closed"      # Normal operation
    OPEN = "open"          # Circuit is open, calls fail fast
    HALF_OPEN = "half_open"  # Testing if service is back


class CircuitBreaker:
    """
    Circuit breaker to prevent cascading failures in exchange connections.
    
    States:
    - CLOSED: Normal operation, requests pass through
    - OPEN: Circuit is open, requests fail immediately
    - HALF_OPEN: Testing if service is back, limited requests allowed
    """
    
    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: int = 60,
        expected_exception: type = Exception,
        name: str = "CircuitBreaker"
    ):
        """
        Initialize circuit breaker.
        
        Args:
            failure_threshold: Number of failures before opening circuit
            recovery_timeout: Time in seconds before attempting recovery
            expected_exception: Exception type that triggers circuit breaker
            name: Name for logging purposes
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.expected_exception = expected_exception
        self.name = name
        
        # State tracking
        self._state = CircuitState.CLOSED
        self._failure_count = 0
        self._last_failure_time: Optional[float] = None
        self._next_attempt_time: Optional[float] = None
        
        logger.info(f"Circuit breaker '{name}' initialized with threshold={failure_threshold}")
    
    @property
    def state(self) -> CircuitState:
        """Get current circuit state"""
        return self._state
    
    @property
    def failure_count(self) -> int:
        """Get current failure count"""
        return self._failure_count
    
    def _should_attempt_reset(self) -> bool:
        """Check if we should attempt to reset the circuit"""
        if self._state != CircuitState.OPEN:
            return False
        
        if self._next_attempt_time is None:
            return False
        
        return time.time() >= self._next_attempt_time
    
    def _on_success(self) -> None:
        """Handle successful operation"""
        if self._state == CircuitState.HALF_OPEN:
            logger.info(f"Circuit breaker '{self.name}' reset to CLOSED after successful test")
            self._state = CircuitState.CLOSED
        
        self._failure_count = 0
        self._last_failure_time = None
        self._next_attempt_time = None
    
    def _on_failure(self) -> None:
        """Handle failed operation"""
        self._failure_count += 1
        self._last_failure_time = time.time()
        
        if self._state == CircuitState.HALF_OPEN:
            # Failed during test, go back to OPEN
            logger.warning(f"Circuit breaker '{self.name}' failed during test, returning to OPEN")
            self._state = CircuitState.OPEN
            self._next_attempt_time = time.time() + self.recovery_timeout
        elif self._failure_count >= self.failure_threshold:
            # Too many failures, open the circuit
            logger.error(
                f"Circuit breaker '{self.name}' OPENED after {self._failure_count} failures"
            )
            self._state = CircuitState.OPEN
            self._next_attempt_time = time.time() + self.recovery_timeout
    
    def call(self, func: Callable, *args, **kwargs) -> Any:
        """
        Execute function with circuit breaker protection.
        
        Args:
            func: Function to execute
            *args: Function arguments
            **kwargs: Function keyword arguments
            
        Returns:
            Function result
            
        Raises:
            CircuitBreakerOpenError: When circuit is open
            Original exception: When function fails
        """
        # Check if we should attempt reset
        if self._should_attempt_reset():
            logger.info(f"Circuit breaker '{self.name}' attempting reset to HALF_OPEN")
            self._state = CircuitState.HALF_OPEN
        
        # Fail fast if circuit is open
        if self._state == CircuitState.OPEN:
            raise CircuitBreakerOpenError(
                f"Circuit breaker '{self.name}' is OPEN. "
                f"Next attempt in {self._next_attempt_time - time.time():.1f}s"
            )
        
        try:
            # Execute the function
            result = func(*args, **kwargs)
            self._on_success()
            return result
            
        except self.expected_exception as e:
            self._on_failure()
            raise e
    
    async def call_async(self, func: Callable, *args, **kwargs) -> Any:
        """
        Execute async function with circuit breaker protection.
        
        Args:
            func: Async function to execute
            *args: Function arguments
            **kwargs: Function keyword arguments
            
        Returns:
            Function result
            
        Raises:
            CircuitBreakerOpenError: When circuit is open
            Original exception: When function fails
        """
        # Check if we should attempt reset
        if self._should_attempt_reset():
            logger.info(f"Circuit breaker '{self.name}' attempting reset to HALF_OPEN")
            self._state = CircuitState.HALF_OPEN
        
        # Fail fast if circuit is open
        if self._state == CircuitState.OPEN:
            raise CircuitBreakerOpenError(
                f"Circuit breaker '{self.name}' is OPEN. "
                f"Next attempt in {self._next_attempt_time - time.time():.1f}s"
            )
        
        try:
            # Execute the async function
            result = await func(*args, **kwargs)
            self._on_success()
            return result
            
        except self.expected_exception as e:
            self._on_failure()
            raise e
    
    def reset(self) -> None:
        """Manually reset the circuit breaker"""
        logger.info(f"Circuit breaker '{self.name}' manually reset")
        self._state = CircuitState.CLOSED
        self._failure_count = 0
        self._last_failure_time = None
        self._next_attempt_time = None
    
    def get_stats(self) -> dict:
        """Get circuit breaker statistics"""
        return {
            'name': self.name,
            'state': self._state.value,
            'failure_count': self._failure_count,
            'failure_threshold': self.failure_threshold,
            'last_failure_time': self._last_failure_time,
            'next_attempt_time': self._next_attempt_time,
            'recovery_timeout': self.recovery_timeout
        }


class CircuitBreakerOpenError(Exception):
    """Exception raised when circuit breaker is open"""
    pass