diff --git a/wizard-work/beam/__pycache__/atomic_write.cpython-314.pyc b/wizard-work/beam/__pycache__/atomic_write.cpython-314.pyc new file mode 100644 index 0000000..962d144 Binary files /dev/null and b/wizard-work/beam/__pycache__/atomic_write.cpython-314.pyc differ diff --git a/wizard-work/beam/test_token_bucket.py b/wizard-work/beam/test_token_bucket.py new file mode 100644 index 0000000..1612d8e --- /dev/null +++ b/wizard-work/beam/test_token_bucket.py @@ -0,0 +1,108 @@ +import threading +import time + +from token_bucket import TokenBucket + + +def test_basic_acquire_and_refill(): + b = TokenBucket(rate=10.0, capacity=5.0) + assert b.try_acquire(5.0) is True + assert b.try_acquire(1.0) is False + # refill: wait ~0.15s -> 1.5 tokens + time.sleep(0.15) + assert b.try_acquire(1.0) is True + + +def test_capacity_cap(): + b = TokenBucket(rate=100.0, capacity=3.0) + time.sleep(0.1) + assert b.tokens() <= 3.0 + 1e-9 + + +def test_acquire_blocks_and_returns(): + b = TokenBucket(rate=20.0, capacity=1.0) + b.try_acquire(1.0) # drain + t0 = time.monotonic() + ok = b.acquire(1.0, timeout=0.2) + dt = time.monotonic() - t0 + assert ok is True + assert dt >= 0.04 # had to wait ~0.05s for 1 token at 20/s + + +def test_n_exceeds_capacity_raises(): + b = TokenBucket(rate=1.0, capacity=1.0) + try: + b.acquire(2.0, timeout=0.01) + raised = False + except ValueError: + raised = True + assert raised + try: + b.try_acquire(2.0) + raised2 = False + except ValueError: + raised2 = True + assert raised2 + + +def test_invalid_params(): + for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}): + try: + TokenBucket(**kwargs) + assert False + except ValueError: + pass + try: + TokenBucket(rate=1, capacity=1).try_acquire(0) + assert False + except ValueError: + pass + + +def test_concurrent_no_overconsume(): + b = TokenBucket(rate=1e-9, capacity=10.0) # effectively no refill + successes = {"n": 0} + lock = threading.Lock() + + def worker(): + if b.try_acquire(1.0): + with lock: + successes["n"] += 1 + + threads = [threading.Thread(target=worker) for _ in range(50)] + for t in threads: + t.start() + for t in threads: + t.join() + assert successes["n"] == 10 + + +def test_acquire_timeout_false(): + b = TokenBucket(rate=1.0, capacity=1.0) + b.try_acquire(1.0) # drain + # request 1 token but timeout too short for refill (need 1s) + t0 = time.monotonic() + ok = b.acquire(1.0, timeout=0.05) + dt = time.monotonic() - t0 + assert ok is False + assert dt < 0.2 + + +def test_reset(): + b = TokenBucket(rate=1.0, capacity=2.0) + b.try_acquire(2.0) + assert b.try_acquire(1.0) is False + b.reset() + assert b.try_acquire(2.0) is True + + +def test_partial_never_consumes(): + b = TokenBucket(rate=1.0, capacity=2.0) + b.try_acquire(2.0) # drain + # 1 token needed, only 0 available -> no consume + assert b.try_acquire(1.0) is False + # refill 0.5s -> 0.5 tokens, still < 1 + time.sleep(0.5) + assert b.try_acquire(1.0) is False + # tokens should be ~0.5, not negative + assert b.tokens() >= 0.4 diff --git a/wizard-work/beam/token_bucket.py b/wizard-work/beam/token_bucket.py new file mode 100644 index 0000000..a0c4686 --- /dev/null +++ b/wizard-work/beam/token_bucket.py @@ -0,0 +1,103 @@ +"""Thread-safe token bucket rate limiter. + +Correctness notes: +- Refill is computed from elapsed wall-clock time using a monotonic clock. +- Tokens are capped at capacity. +- acquire(n) blocks until n tokens available or timeout; partial consumes never happen. +- A single lock guards state; condition variable for waiters. +""" +from __future__ import annotations + +import threading +import time +from typing import Optional + + +class TokenBucket: + def __init__(self, rate: float, capacity: float): + if rate <= 0: + raise ValueError("rate must be > 0") + if capacity <= 0: + raise ValueError("capacity must be > 0") + if rate > capacity: + # allow; but typically rate <= capacity. No restriction. + pass + self._rate = float(rate) + self._capacity = float(capacity) + self._tokens = float(capacity) + self._last = time.monotonic() + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + @property + def rate(self) -> float: + return self._rate + + @property + def capacity(self) -> float: + return self._capacity + + def _refill(self, now: float) -> None: + elapsed = now - self._last + if elapsed <= 0: + return + self._tokens = min(self._capacity, self._tokens + elapsed * self._rate) + self._last = now + + def tokens(self) -> float: + with self._lock: + self._refill(time.monotonic()) + return self._tokens + + def try_acquire(self, n: float = 1.0) -> bool: + if n <= 0: + raise ValueError("n must be > 0") + if n > self._capacity: + raise ValueError("n exceeds capacity") + with self._lock: + now = time.monotonic() + self._refill(now) + if self._tokens >= n - 1e-12: + self._tokens -= n + return True + return False + + def acquire(self, n: float = 1.0, timeout: Optional[float] = None) -> bool: + if n <= 0: + raise ValueError("n must be > 0") + if n > self._capacity: + raise ValueError("n exceeds capacity") + deadline = None if timeout is None else time.monotonic() + timeout + with self._lock: + while True: + now = time.monotonic() + self._refill(now) + if self._tokens >= n - 1e-12: + self._tokens -= n + return True + # time until enough tokens + deficit = n - self._tokens + wait = deficit / self._rate + if deadline is not None: + remaining = deadline - now + if remaining <= 0: + return False + wait = min(wait, remaining) + # Guard against sub-millisecond waits that degenerate into a + # busy spin on spurious wakeups (cond.wait(1e-15) returns at + # once). Never wait longer than the allowed remaining time. + if wait < 1e-3: + if deadline is not None: + if remaining <= 1e-3: + return False + wait = min(remaining, 1e-3) + else: + wait = 1e-3 + self._cond.wait(timeout=wait) + # loop rechecks + + def reset(self, tokens: Optional[float] = None) -> None: + with self._lock: + self._tokens = self._capacity if tokens is None else float(tokens) + self._last = time.monotonic() + self._cond.notify_all()