Antoni Nuñez Romeu
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
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,
|
|
}
|