Circuit Breaker

circuit-breaker.py

"""
Circuit Breaker Imnplementation.
"""

from datetime import datetime

# constants
CLOSED = 0
OPEN = 1
HALF_OPEN = 2
DEFAULT_THRESHOLD = 5  # failure count threshhold
DEFAULT_TIMEOUT = 60  # threshhold timeout in seconds

class FailureException(Exception):
    pass

class CircuitBreakerOpen(Exception):
    pass

class CircuitBreaker:

    def __init__(self, *args, **kwargs):
        self.state = CLOSED
        self.threshold = kwargs.get('threshold', DEFAULT_THRESHOLD)
        self.timeout = kwargs.get('timeout', DEFAULT_TIMEOUT)
        self.failures = 0
        self.last_failed = None

    def _set_state(self):
        if self.failures > self.threshold:
            timeout = datetime.now() - self.last_failed
            if timeout.total_seconds() > self.timeout:
                self.state = HALF_OPEN
            else:
                self.state = OPEN
        else:
            self.state = CLOSED

    def _reset_state(self):
        self.failures = 0
        self.last_failed = None
        self.state = CLOSED
    
    def _set_failures(self):
        self.failures += 1
        self.last_failed = datetime.now()

    def run(self, callable):
        self._set_state()
        if self.state in (CLOSED, HALF_OPEN) :
            try:
                return callable()
            except Exception as e:
                self._set_failures()
                raise FailureException("Failure!")
            else:
                self._reset_state()
        elif self.state == OPEN:
            raise CircuitBreakerOpen("Circuit Open!")