Chaos testing an ML service

import random


class CircuitBreaker:
    def __init__(self, failure_threshold=3, cooldown_calls=3):
        self.failure_threshold = failure_threshold
        self.cooldown_calls = cooldown_calls
        self.state = "CLOSED"
        self.consecutive_failures = 0
        self.calls_since_opened = 0

    def call(self, provider_fn, fallback_fn, *args):
        if self.state == "OPEN":
            self.calls_since_opened += 1
            if self.calls_since_opened >= self.cooldown_calls:
                self.state = "HALF_OPEN"
            else:
                return fallback_fn(*args)
        try:
            result = provider_fn(*args)
        except RuntimeError:
            self.consecutive_failures += 1
            if self.state == "HALF_OPEN" or self.consecutive_failures >= self.failure_threshold:
                self.state = "OPEN"
                self.calls_since_opened = 0
            return fallback_fn(*args)
        self.state = "CLOSED"
        self.consecutive_failures = 0
        return result


def run_chaos_experiment(cache: dict, seed: int, n_requests: int = 50, failure_rate: float = 0.35):
    """A fixed seed makes this reproducible. Real chaos tooling injects
    failure into a live system; this injects into a local simulation of
    one, which is what a single laptop can honestly demonstrate."""
    rng = random.Random(seed)
    cities = ["MUM", "DEL", "BLR", "PUN", "HYD", "KOL"]

    def flaky_provider(city):
        if rng.random() < failure_rate:
            raise RuntimeError("provider timeout")
        return f"live({city})"

    def cached_fallback(city):
        if city not in cache:
            raise KeyError(f"no cached value for {city}")
        return f"cached({city})"

    breaker = CircuitBreaker()
    ok, unhandled = 0, []
    for _ in range(n_requests):
        city = rng.choice(cities)
        try:
            breaker.call(flaky_provider, cached_fallback, city)
            ok += 1
        except KeyError as e:
            unhandled.append(str(e))
    return ok, unhandled


# The cache only ever warmed up for the three cities the system has always
# served. Pune, Hyderabad and Kolkata are newer and have no cache entry.
warm_cache = {"MUM": 1, "DEL": 1, "BLR": 1}
ok, unhandled = run_chaos_experiment(warm_cache, seed=7)
print("Experiment 1: cold cache for new cities")
print(f"  {ok} requests survived, {len(unhandled)} were unhandled failures")
for u in unhandled[:3]:
    print(f"    unhandled: {u}")
if len(unhandled) > 3:
    print(f"    ... and {len(unhandled) - 3} more")

from chaos_experiment import CircuitBreaker
import random

warm_cache = {"MUM": 1, "DEL": 1, "BLR": 1}


def cached_fallback_with_default(cache, city):
    if city in cache:
        return f"cached({city})"
    return "default-neutral-score"


def run_chaos_experiment_fixed(cache: dict, seed: int, n_requests: int = 50, failure_rate: float = 0.35):
    rng = random.Random(seed)
    cities = ["MUM", "DEL", "BLR", "PUN", "HYD", "KOL"]

    def flaky_provider(city):
        if rng.random() < failure_rate:
            raise RuntimeError("provider timeout")
        return f"live({city})"

    breaker = CircuitBreaker()
    ok, unhandled = 0, []
    for _ in range(n_requests):
        city = rng.choice(cities)
        try:
            breaker.call(flaky_provider, lambda c: cached_fallback_with_default(cache, c), city)
            ok += 1
        except Exception as e:
            unhandled.append(str(e))
    return ok, unhandled


ok2, unhandled2 = run_chaos_experiment_fixed(warm_cache, seed=7)
print("Experiment 2: same chaos, with a final default answer added")
print(f"  {ok2} requests survived, {len(unhandled2)} were unhandled failures")