diff --git a/wizard-work/beam/retry_jitter.py b/wizard-work/beam/retry_jitter.py new file mode 100644 index 0000000..4b62b00 --- /dev/null +++ b/wizard-work/beam/retry_jitter.py @@ -0,0 +1,39 @@ +"""Bounded retry with exponential backoff + full jitter. + +Fixes subtle bug in prior retry: jitter must be uniform over [0, cap] where cap +is min(max_delay, base_delay*2**attempt), and a zero/negative sleep must never +be passed to a clock that rejects negatives. Also re-raises the ORIGINAL +exception instance (with its traceback) rather than a re-wrapped one. +""" +from __future__ import annotations +import random, time +from typing import Callable, Optional, Tuple, Type + + +class RetryExhausted(Exception): + pass + + +def retry(fn, *, max_attempts=3, base_delay=0.1, max_delay=10.0, + retryable=(Exception,), sleep=time.sleep, rng=None): + 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 = None + for attempt in range(max_attempts): + try: + return fn() + except retryable as e: + last = e + if attempt == max_attempts - 1: + break + exp = min(max_delay, base_delay * (2 ** attempt)) + delay = r.uniform(0.0, exp) if exp > 0 else 0.0 + if delay < 0: + delay = 0.0 + sleep(delay) + if last is None: + raise RetryExhausted() + raise last diff --git a/wizard-work/beam/test_retry_jitter.py b/wizard-work/beam/test_retry_jitter.py new file mode 100644 index 0000000..9634302 --- /dev/null +++ b/wizard-work/beam/test_retry_jitter.py @@ -0,0 +1,64 @@ +import sys, os, time +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from retry_jitter import retry + +def test_success_no_sleep(): + calls = {"n":0} + sleeps = [] + def fn(): + calls["n"]+=1 + return 42 + assert retry(fn, sleep=sleeps.append) == 42 + assert calls["n"]==1 and sleeps==[] + +def test_retries_then_succeeds(): + calls={"n":0} + def fn(): + calls["n"]+=1 + if calls["n"]<3: raise ValueError("boom") + return "ok" + assert retry(fn, max_attempts=3, base_delay=0, sleep=lambda x:None)=="ok" + assert calls["n"]==3 + +def test_exhausted_raises_original(): + class E(Exception): pass + def fn(): + raise E("x") + try: + retry(fn, max_attempts=2, base_delay=0, sleep=lambda x:None) + except E as e: + assert str(e)=="x" + else: + raise AssertionError("no raise") + +def test_nonretryable_propagates(): + class A(Exception): pass + class B(Exception): pass + def fn(): + raise B("no") + try: + retry(fn, retryable=(A,), max_attempts=3, base_delay=0, sleep=lambda x:None) + except B: + pass + else: + raise AssertionError("not propagated") + +def test_max_attempts_validation(): + try: + retry(lambda: 1, max_attempts=0) + except ValueError: + pass + else: + raise AssertionError("no validation") + +def test_jitter_bounded(): + rng=__import__("random").Random(1) + sleeps=[] + def fn(): + raise ValueError("e") + try: + retry(fn, max_attempts=4, base_delay=1.0, max_delay=2.0, rng=rng, sleep=sleeps.append) + except ValueError: + pass + for d in sleeps: + assert 0.0 <= d <= 2.0