from __future__ import annotations import time from enum import Enum from typing import Any, Callable from src.core.exceptions import ConnectorException, ServiceUnavailableException from src.core.logging import get_logger logger = get_logger("circuit_breaker") class CircuitState(str, Enum): CLOSED = "closed" OPEN = "open" HALF_OPEN = "half_open" class CircuitBreaker: def __init__( self, name: str, failure_threshold: int = 5, recovery_timeout: float = 30.0, half_open_max_calls: int = 1, ) -> None: self._name = name self._failure_threshold = failure_threshold self._recovery_timeout = recovery_timeout self._half_open_max_calls = half_open_max_calls self._state = CircuitState.CLOSED self._failure_count = 0 self._success_count = 0 self._last_failure_time: float = 0 self._half_open_calls: int = 0 @property def state(self) -> CircuitState: if self._state == CircuitState.OPEN: if time.monotonic() - self._last_failure_time >= self._recovery_timeout: self._state = CircuitState.HALF_OPEN self._half_open_calls = 0 logger.info("circuit_half_open", extra={"name": self._name}) return self._state async def call(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: current_state = self.state if current_state == CircuitState.OPEN: raise ServiceUnavailableException( f"circuit_breaker:{self._name}", f"Circuit breaker '{self._name}' is open. Retry after {self._recovery_timeout}s.", ) if current_state == CircuitState.HALF_OPEN: if self._half_open_calls >= self._half_open_max_calls: raise ServiceUnavailableException( f"circuit_breaker:{self._name}", f"Circuit breaker '{self._name}' is half-open with max concurrent calls reached.", ) self._half_open_calls += 1 try: result = await func(*args, **kwargs) self._on_success() return result except Exception as exc: self._on_failure(exc) raise def _on_success(self) -> None: if self._state == CircuitState.HALF_OPEN: self._state = CircuitState.CLOSED logger.info("circuit_closed", extra={"name": self._name}) self._failure_count = 0 self._success_count += 1 def _on_failure(self, exc: Exception) -> None: self._failure_count += 1 self._last_failure_time = time.monotonic() self._success_count = 0 if self._state == CircuitState.HALF_OPEN: self._state = CircuitState.OPEN logger.warning("circuit_reopened", extra={"name": self._name, "error": str(exc)}) elif self._failure_count >= self._failure_threshold: self._state = CircuitState.OPEN logger.warning("circuit_opened", extra={"name": self._name, "failure_count": self._failure_count}) def reset(self) -> None: self._state = CircuitState.CLOSED self._failure_count = 0 self._success_count = 0 self._last_failure_time = 0 def get_status(self) -> dict[str, Any]: return { "name": self._name, "state": self.state.value, "failure_count": self._failure_count, "success_count": self._success_count, "failure_threshold": self._failure_threshold, "recovery_timeout": self._recovery_timeout, "last_failure_time": self._last_failure_time, }