beam: fix token_bucket acquire busy-spin on sub-ms waits

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

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()