beam: add retry_jitter with bounded full-jitter backoff and tests

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

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