beam: add bounded retry with exponential backoff + tests

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

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