beam: token-bucket rate limiter with injectable clock + tests

wizard-work/beam/ratelimit.py

@@ -0,0 +1,87 @@
+"""Token-bucket rate limiter with injectable clock and thread safety."""
+
+from __future__ import annotations
+
+import threading
+import time
+from dataclasses import dataclass, field
+
+
+class RateLimitError(ValueError):
+    pass
+
+
+@dataclass
+class TokenBucket:
+    rate: float
+    capacity: float
+    _tokens: float = field(init=False, repr=False)
+    _last: float = field(init=False, repr=False)
+    _lock: threading.Lock = field(init=False, repr=False, default_factory=threading.Lock)
+    _clock: object = field(init=False, repr=False, default=None)
+
+    def __post_init__(self) -> None:
+        if self.rate <= 0:
+            raise ValueError("rate must be > 0")
+        if self.capacity <= 0:
+            raise ValueError("capacity must be > 0")
+        if self._clock is None:
+            self._clock = time.monotonic
+        self._tokens = self.capacity
+        self._last = self._clock()
+
+    def tokens(self) -> float:
+        with self._lock:
+            self._refill()
+            return self._tokens
+
+    def acquire(self, n: float = 1.0, *, blocking: bool = True,
+                timeout: float | None = None) -> bool:
+        if n < 0:
+            raise ValueError("n must be >= 0")
+        if n == 0:
+            return True
+        if n > self.capacity and not blocking:
+            return False
+        if n > self.capacity:
+            raise RateLimitError(f"request {n} exceeds capacity {self.capacity}")
+
+        deadline = None
+        if timeout is not None:
+            deadline = self._clock() + timeout
+
+        while True:
+            with self._lock:
+                self._refill()
+                if self._tokens >= n:
+                    self._tokens -= n
+                    return True
+                deficit = n - self._tokens
+                wait = deficit / self.rate
+                now = self._clock()
+                if deadline is not None and now + wait > deadline:
+                    return False
+                if not blocking:
+                    return False
+                if deadline is not None:
+                    wait = min(wait, deadline - now)
+                sleep_for = wait
+            if sleep_for > 0:
+                time.sleep(sleep_for)
+
+    def _refill(self) -> None:
+        now = self._clock()
+        elapsed = now - self._last
+        if elapsed < 0:
+            elapsed = 0.0
+        self._tokens = min(self.capacity, self._tokens + elapsed * self.rate)
+        self._last = now
+
+
+def make_bucket(rate, capacity, clock=None):
+    b = TokenBucket(rate=rate, capacity=capacity)
+    if clock is not None:
+        b._clock = clock
+        b._tokens = capacity
+        b._last = clock()
+    return b

wizard-work/beam/test_ratelimit.py

@@ -0,0 +1,96 @@
+import time
+import threading
+import pytest
+from ratelimit import TokenBucket, make_bucket, RateLimitError
+
+
+def test_basic_burst():
+    b = make_bucket(rate=10, capacity=5, clock=lambda: 0.0)
+    # capacity 5, all available at t=0
+    for _ in range(5):
+        assert b.acquire(1, blocking=False) is True
+    assert b.acquire(1, blocking=False) is False
+
+
+def test_refill():
+    t = [0.0]
+    b = make_bucket(rate=2, capacity=4, clock=lambda: t[0])
+    assert b.acquire(4, blocking=False) is True
+    assert b.acquire(1, blocking=False) is False
+    t[0] = 1.0  # 2 tokens refilled
+    assert b.acquire(2, blocking=False) is True
+    assert b.acquire(1, blocking=False) is False
+
+
+def test_acquire_zero():
+    b = make_bucket(rate=1, capacity=1, clock=lambda: 0.0)
+    assert b.acquire(0) is True
+    assert b.tokens() == 1.0
+
+
+def test_capacity_cap():
+    t = [0.0]
+    b = make_bucket(rate=10, capacity=3, clock=lambda: t[0])
+    t[0] = 100.0
+    assert b.tokens() == 3.0  # capped, not 1003
+
+
+def test_timeout():
+    t = [0.0]
+    def clock():
+        return t[0]
+    b = make_bucket(rate=1, capacity=1, clock=clock)
+    assert b.acquire(1, blocking=False) is True
+    # need 1 more token, refill 1/s; timeout 0.5 -> fail
+    ok = b.acquire(1, blocking=True, timeout=0.5)
+    assert ok is False
+
+
+def test_blocking_success_with_fake_sleep(monkeypatch):
+    t = [0.0]
+    b = make_bucket(rate=2, capacity=2, clock=lambda: t[0])
+    assert b.acquire(2, blocking=False) is True
+    # monkeypatch time.sleep to advance clock
+    def fake_sleep(s):
+        t[0] += s
+    monkeypatch.setattr(time, "sleep", fake_sleep)
+    assert b.acquire(2, blocking=True, timeout=5) is True
+
+
+def test_n_exceeds_capacity_nonblocking():
+    b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
+    assert b.acquire(5, blocking=False) is False
+
+
+def test_n_exceeds_capacity_blocking_raises():
+    b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
+    with pytest.raises(RateLimitError):
+        b.acquire(5, blocking=True)
+
+
+def test_negative_n():
+    b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
+    with pytest.raises(ValueError):
+        b.acquire(-1, blocking=False)
+
+
+def test_thread_safety():
+    b = make_bucket(rate=1000, capacity=100, clock=time.monotonic)
+    results = []
+    def worker():
+        for _ in range(50):
+            results.append(b.acquire(1, blocking=True))
+    threads = [threading.Thread(target=worker) for _ in range(4)]
+    for th in threads:
+        th.start()
+    for th in threads:
+        th.join()
+    assert all(results)
+    assert len(results) == 200
+
+
+def test_invalid_params():
+    with pytest.raises(ValueError):
+        TokenBucket(rate=0, capacity=1)
+    with pytest.raises(ValueError):
+        TokenBucket(rate=1, capacity=0)