diff --git a/wizard-work/beam/ratelimit.py b/wizard-work/beam/ratelimit.py new file mode 100644 index 0000000..26d1867 --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_ratelimit.py b/wizard-work/beam/test_ratelimit.py new file mode 100644 index 0000000..70b318a --- /dev/null +++ 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)