120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
"""Thread-safe circuit breaker.
|
|
|
|
States: CLOSED (normal), OPEN (failing fast), HALF_OPEN (probing).
|
|
Transitions:
|
|
CLOSED -> OPEN after `failure_threshold` consecutive failures.
|
|
OPEN -> HALF_OPEN after `recovery_timeout` elapsed.
|
|
HALF_OPEN -> CLOSED after `success_threshold` consecutive successes.
|
|
HALF_OPEN -> OPEN on any failure.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from enum import Enum
|
|
from typing import Any, Callable, Optional
|
|
|
|
|
|
class CircuitState(Enum):
|
|
CLOSED = "closed"
|
|
OPEN = "open"
|
|
HALF_OPEN = "half_open"
|
|
|
|
|
|
class CircuitOpenError(Exception):
|
|
"""Raised when a call is rejected because the circuit is open."""
|
|
|
|
|
|
class CircuitBreaker:
|
|
def __init__(
|
|
self,
|
|
failure_threshold: int = 5,
|
|
recovery_timeout: float = 30.0,
|
|
success_threshold: int = 2,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
):
|
|
if failure_threshold < 1:
|
|
raise ValueError("failure_threshold must be >= 1")
|
|
if recovery_timeout < 0:
|
|
raise ValueError("recovery_timeout must be >= 0")
|
|
if success_threshold < 1:
|
|
raise ValueError("success_threshold must be >= 1")
|
|
self._failure_threshold = failure_threshold
|
|
self._recovery_timeout = recovery_timeout
|
|
self._success_threshold = success_threshold
|
|
self._clock = clock
|
|
self._state = CircuitState.CLOSED
|
|
self._consecutive_failures = 0
|
|
self._consecutive_successes = 0
|
|
self._opened_at: Optional[float] = None
|
|
self._probe_in_flight = False
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def state(self) -> CircuitState:
|
|
with self._lock:
|
|
self._maybe_transition_to_half_open()
|
|
return self._state
|
|
|
|
def _maybe_transition_to_half_open(self) -> None:
|
|
if self._state is CircuitState.OPEN and self._opened_at is not None:
|
|
if self._clock() - self._opened_at >= self._recovery_timeout:
|
|
self._state = CircuitState.HALF_OPEN
|
|
self._consecutive_successes = 0
|
|
|
|
def _on_success(self) -> None:
|
|
with self._lock:
|
|
self._maybe_transition_to_half_open()
|
|
if self._state is CircuitState.HALF_OPEN:
|
|
self._consecutive_successes += 1
|
|
if self._consecutive_successes >= self._success_threshold:
|
|
self._state = CircuitState.CLOSED
|
|
self._consecutive_failures = 0
|
|
self._consecutive_successes = 0
|
|
self._opened_at = None
|
|
else:
|
|
self._consecutive_failures = 0
|
|
|
|
def _on_failure(self) -> None:
|
|
with self._lock:
|
|
self._maybe_transition_to_half_open()
|
|
if self._state is CircuitState.HALF_OPEN:
|
|
self._state = CircuitState.OPEN
|
|
self._opened_at = self._clock()
|
|
self._consecutive_successes = 0
|
|
self._consecutive_failures = 0
|
|
return
|
|
self._consecutive_failures += 1
|
|
if self._consecutive_failures >= self._failure_threshold:
|
|
self._state = CircuitState.OPEN
|
|
self._opened_at = self._clock()
|
|
|
|
def call(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
|
# Check gate.
|
|
with self._lock:
|
|
self._maybe_transition_to_half_open()
|
|
if self._state is CircuitState.OPEN:
|
|
raise CircuitOpenError("circuit is open")
|
|
if self._state is CircuitState.HALF_OPEN:
|
|
if self._probe_in_flight:
|
|
raise CircuitOpenError("circuit is open (probe in flight)")
|
|
self._probe_in_flight = True
|
|
try:
|
|
result = fn(*args, **kwargs)
|
|
except BaseException:
|
|
self._on_failure()
|
|
raise
|
|
else:
|
|
self._on_success()
|
|
return result
|
|
finally:
|
|
with self._lock:
|
|
self._probe_in_flight = False
|
|
|
|
def reset(self) -> None:
|
|
with self._lock:
|
|
self._state = CircuitState.CLOSED
|
|
self._consecutive_failures = 0
|
|
self._consecutive_successes = 0
|
|
self._opened_at = None
|
|
self._probe_in_flight = False |