diff --git a/wizard-work/beam/__pycache__/retry.cpython-314.pyc b/wizard-work/beam/__pycache__/retry.cpython-314.pyc new file mode 100644 index 0000000..eb930ae Binary files /dev/null and b/wizard-work/beam/__pycache__/retry.cpython-314.pyc differ diff --git a/wizard-work/beam/__pycache__/test_atomic_write.cpython-314.pyc b/wizard-work/beam/__pycache__/test_atomic_write.cpython-314.pyc new file mode 100644 index 0000000..6b0cf67 Binary files /dev/null and b/wizard-work/beam/__pycache__/test_atomic_write.cpython-314.pyc differ diff --git a/wizard-work/beam/__pycache__/test_retry.cpython-314.pyc b/wizard-work/beam/__pycache__/test_retry.cpython-314.pyc new file mode 100644 index 0000000..5712486 Binary files /dev/null and b/wizard-work/beam/__pycache__/test_retry.cpython-314.pyc differ diff --git a/wizard-work/beam/__pycache__/test_token_bucket.cpython-314.pyc b/wizard-work/beam/__pycache__/test_token_bucket.cpython-314.pyc new file mode 100644 index 0000000..42ba3a0 Binary files /dev/null and b/wizard-work/beam/__pycache__/test_token_bucket.cpython-314.pyc differ diff --git a/wizard-work/beam/__pycache__/token_bucket.cpython-314.pyc b/wizard-work/beam/__pycache__/token_bucket.cpython-314.pyc new file mode 100644 index 0000000..9386298 Binary files /dev/null and b/wizard-work/beam/__pycache__/token_bucket.cpython-314.pyc differ diff --git a/wizard-work/beam/retry.py b/wizard-work/beam/retry.py new file mode 100644 index 0000000..dc1ab7b --- /dev/null +++ b/wizard-work/beam/retry.py @@ -0,0 +1,38 @@ +"""Bounded retry with exponential backoff + jitter.""" +from __future__ import annotations + +import random +import time +from typing import Callable, Optional, Tuple, Type + + +def retry( + fn: Callable, + *, + max_attempts: int = 3, + base_delay: float = 0.1, + max_delay: float = 10.0, + retryable: Tuple[Type[BaseException], ...] = (Exception,), + sleep: Callable[[float], None] = time.sleep, + rng: Optional[random.Random] = None, +) -> object: + if max_attempts < 1: + raise ValueError("max_attempts must be >= 1") + if base_delay < 0 or max_delay < 0: + raise ValueError("delays must be >= 0") + r = rng if rng is not None else random.Random() + last_exc = None + for attempt in range(max_attempts): + try: + return fn() + except retryable as e: + last_exc = e + if attempt == max_attempts - 1: + break + exp = min(max_delay, base_delay * (2 ** attempt)) + delay = r.uniform(0, exp) if exp > 0 else 0.0 + sleep(delay) + except BaseException: + raise + assert last_exc is not None + raise last_exc diff --git a/wizard-work/beam/test_retry.py b/wizard-work/beam/test_retry.py new file mode 100644 index 0000000..03f13e9 --- /dev/null +++ b/wizard-work/beam/test_retry.py @@ -0,0 +1,84 @@ +import time +from retry import retry + + +def test_succeeds_first(): + calls = {"n": 0} + def f(): + calls["n"] += 1 + return "ok" + assert retry(f) == "ok" + assert calls["n"] == 1 + + +def test_retries_then_succeeds(): + calls = {"n": 0} + def f(): + calls["n"] += 1 + if calls["n"] < 3: + raise ValueError("nope") + return "done" + assert retry(f, max_attempts=5, base_delay=0, max_delay=0) == "done" + assert calls["n"] == 3 + + +def test_exhausts_and_reraises(): + calls = {"n": 0} + def f(): + calls["n"] += 1 + raise RuntimeError("boom") + try: + retry(f, max_attempts=2, base_delay=0, max_delay=0) + assert False + except RuntimeError as e: + assert str(e) == "boom" + assert calls["n"] == 2 + + +def test_non_retryable_propagates_immediately(): + calls = {"n": 0} + def f(): + calls["n"] += 1 + raise KeyError("x") + try: + retry(f, max_attempts=5, base_delay=0, retryable=(ValueError,)) + assert False + except KeyError: + pass + assert calls["n"] == 1 + + +def test_invalid_params(): + for kw in ({"max_attempts": 0}, {"base_delay": -1}, {"max_delay": -1}): + try: + retry(lambda: None, **kw) + assert False + except ValueError: + pass + + +def test_sleep_called_with_bounded_delay(): + sleeps = [] + calls = {"n": 0} + def f(): + calls["n"] += 1 + if calls["n"] < 3: + raise ValueError("x") + return "y" + retry(f, max_attempts=3, base_delay=1.0, max_delay=2.0, sleep=sleeps.append) + assert len(sleeps) == 2 + for d in sleeps: + assert 0 <= d <= 2.0 + + +def test_keyboard_interrupt_not_retried(): + calls = {"n": 0} + def f(): + calls["n"] += 1 + raise KeyboardInterrupt() + try: + retry(f, max_attempts=5, base_delay=0) + assert False + except KeyboardInterrupt: + pass + assert calls["n"] == 1