熔断器阻止对故障服务的轰炸

from concurrent.futures import ThreadPoolExecutor
from threading import Barrier, Lock


class OpenCircuit(Exception):
    pass


class Breaker:
    # Explicit fixture policy: two consecutive classified failures, five-second cooldown.
    def __init__(self):
        self.state, self.failures, self.until = "closed", 0, 0
        self.epoch, self.probe_busy = 0, False
        self.lock = Lock()

    def acquire(self, now):
        with self.lock:
            if self.state == "open":
                if now < self.until:
                    raise OpenCircuit
                self.state = "half-open"
            if self.state == "half-open":
                if self.probe_busy:
                    raise OpenCircuit
                self.probe_busy = True
                return self.epoch, True
            return self.epoch, False

    def complete(self, permit, succeeded, now):
        epoch, probe = permit
        with self.lock:
            if epoch != self.epoch:
                return  # This outcome belongs to an earlier breaker generation.
            if probe:
                self.probe_busy = False
                self.epoch += 1
                self.failures = 0
                self.state = "closed" if succeeded else "open"
                if not succeeded:
                    self.until = now + 5
            elif self.state == "closed":
                self.failures = 0 if succeeded else self.failures + 1
                if self.failures >= 2:
                    self.state, self.until = "open", now + 5
                    self.epoch += 1


breaker = Breaker()
other_dependency = Breaker()
forwarded = []


def invoke(succeeded, now):
    permit = breaker.acquire(now)
    forwarded.append(now)  # A modeled dependency invocation, not a network call.
    breaker.complete(permit, succeeded, now)


invoke(False, 0)
assert breaker.failures == 1
invoke(True, 1)
assert breaker.failures == 0
late = breaker.acquire(1)
invoke(False, 2)
invoke(False, 3)
assert breaker.state == "open" and breaker.until == 8
other_permit = other_dependency.acquire(4)
other_dependency.complete(other_permit, True, 4)
assert other_dependency.state == "closed"
before = len(forwarded)
try:
    invoke(True, 7.99)
except OpenCircuit:
    pass
else:
    raise AssertionError("open circuit must reject before dispatch")
assert len(forwarded) == before
breaker.complete(late, True, 7.99)
assert breaker.state == "open"
print("success resets consecutive failures; two later failures open until time 8")
print("open calls are not dispatched; an older in-flight success does not close a newer open generation")

gate = Barrier(8)


def probe_attempt(_):
    gate.wait(timeout=5)
    try:
        return breaker.acquire(8)
    except OpenCircuit:
        return None


with ThreadPoolExecutor(max_workers=8) as pool:
    permits = list(pool.map(probe_attempt, range(8)))
admitted = [permit for permit in permits if permit is not None]
assert len(admitted) == 1 and admitted[0][1] is True
assert breaker.state == "half-open" and breaker.probe_busy
breaker.complete(admitted[0], False, 8)
assert breaker.state == "open" and breaker.until == 13
print("eight simultaneous half-open candidates receive one probe permit and seven refusals")
print("failed probe opens a new cooldown ending at time 13")

try:
    breaker.acquire(12.99)
except OpenCircuit:
    pass
else:
    raise AssertionError("second cooldown must still apply")
probe = breaker.acquire(13)
breaker.complete(probe, True, 13)
assert breaker.state == "closed" and not breaker.probe_busy
invoke(True, 14)
assert other_dependency.state == "closed" and other_dependency.failures == 0
print("a successful later probe closes this breaker; another dependency has separate state")