From 84c7c1564443bb847d09e780d8d64398f06ec687 Mon Sep 17 00:00:00 2001 From: beam Date: Tue, 20 Oct 2026 19:54:19 +0000 Subject: [PATCH] beam: fix token_bucket acquire busy-spin on sub-ms waits --- .../__pycache__/atomic_write.cpython-314.pyc | Bin 0 -> 5503 bytes wizard-work/beam/test_token_bucket.py | 108 ++++++++++++++++++ wizard-work/beam/token_bucket.py | 103 +++++++++++++++++ 3 files changed, 211 insertions(+) create mode 100644 wizard-work/beam/__pycache__/atomic_write.cpython-314.pyc create mode 100644 wizard-work/beam/test_token_bucket.py create mode 100644 wizard-work/beam/token_bucket.py 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 0000000000000000000000000000000000000000..962d144211237c3d1e94c2fa9fc45c7ee1cd8380 GIT binary patch literal 5503 zcmcIoYitzP6~6P>ot=GqSw8^7jvv8Z@S~z8BrypXLhzbmyA!iC31l_ioy805oz0!u z`0*ohqgEv-Qk*CcXd*$RN`b0UQI#M3s4Ch<3J|Gsc*yS5RicJJ`L_t#L`kLe+&eS7 zi*b{vQm?&p&)j?Nx#ygF?mgd~Ex~{nL3#1li^u-pL+CT|;uO&$m^<$PF^6Pipl6ZH zB$z?gU}J=46P&@(nm2f;xrA%bZMbP4pAZH;hG$ST#6hp&jiDoM*(JMWLH5X^?3I18 zx0jdwjYtkOqI@{d_Sy3qK8L55m4nb9vic1_8Q06mVQ8zf+X6D%hT@UZ&sl^93~e%( z?vm12(I>ifC96u=Twax?aL&jTj`n!Hhk#Ob$lyhJuf(hSaHRn$StKO)C_p#pspNMAMr|1vRIq6 z)kqPyfr-RR#yxG6b1DZ0SqfX6ovm8Q=rN{^);Mum))GjC_E1rstrKcKueHYM_%_?^ zx|AtlWgM1ii!F}$601fD7rfqqJUCo06^RbIG^u6OM_ZLvsZ%Ky)j}pNjT28P8CYF# zOHL=tETy%Ap%#pt(y%7|WA1gy;swjS!y{wP9QJb4_0rXXuJ%hMUA1SJYL&7Yma?UM z9>m+;Kcq^>RkbK7UL9rueTg=dFRft`Ed^EAd%SVR^eKgcW+-Hd^kH}pQjMmKNu>ae z#&OPM6CG|Vh$zIv+P4ZLns4n&FII>1S2ytD|7JUD;rUuTgnYv$oDzy zAY}4I#h5TT-N0CYcTCQJ*?^3)xAkI8GkOiIm5OS&Za|clioH|0Sp{dhr!;)LcU)B_duxNLr#NH!Qz_e_ zYAS{6LDYKq=}*984t;^PGxfL0W62vRb5NJv$QLXNgz}(Jb6F&m7x}|wAEEpx+Efk@ zDu~>X+aaLtfFwj+NeCW(#OgMLTz6PED2`&pos1lK80^T@5qBPToDSz$6)Yn8U=Gw2TsYE6ONw?TYdpBkZ%dc2E*I^Eq`t(_u3lrHdYe6S?TBkxb9}S{Vqn zprnfg#CBU_7~?Lwle{A-bAa0W?e+xjl}X;J+EpZVu8=o*#p*zj%`z51)g4yl5Vweg)R-%Q4ZCD4n5yj@k)&sn;ed6-GZKa zPWC1l8X?s8W9$J2vK=q`w%ZX<6y;!1fZ4@qG{?lC9wmXiqBvO)`nihd9$2Am^NK<6 z3Z4^3WB0c^@|+l<^S1nVz8I>pwDo>I@XB1(D-bK58aMeq7ry;t{-x9Wqc?AJDTaenVvDoo*oKm zQrhw1+&2uQVNZ{p8sn_#1{{#hO`9A*NRt_lvouqiLOQSMYBpo~MpBO_K{qJ=Fg*l*60=QV@;E`fMacQ-2`~)mj0PYLkW0R#PgtPAO;`(>>7r_yf(@=r zo^pCdQu5;?!#{e;6aZ^0hJnq%Zbeu3I)K*{tE9VB$ma^jO=09oYDU~gji{m)I->

0UoB>)hS~ zczWYGesOL0wYA;LP3z7GzZA~$ZwD5ex<78}zTmpqwCzOTR8Ud>oPOb4dNH>3T5Kx}4Ac{S#Mi`AZEpr| zyO8w5mv(-6w+_`eTGHMa`Fc6jd>6UI2<2W1*O7I2T?F3-B46Hhq3DLYhzZm$g`-RU zy0Qz%ugU_je|VP1l>;T{I~uW2%8rN z*6rWP{Bc7&JpaVfXU6~oFMp0dHq?PW=-j;}(S|;ZHY7ag!_5sqUq&o+TxL6g>2g=Y zP(8ZBx`DppX-Y`k6^R3$D{a1^2zRA7JcQUQI~ky_BHvKZb=AibI=~P*$P+rkfwZgj zzM*#Rs>Feut8EOS+bP`<-@k+VDAKS$&VSU&5IXMLzg_sKk0tbj457EX2)%>w-*X2c z029(QbBAP%@3@1o&MlnQZlL;Ec0KTDTSSzTQAip?+ z*Fa_6FX&FC_XGMkNIFP5B9{b~X}C=ur}mUdbyla<#Wl3LY|v^Um8LwHV8OHv{1B-L z;GkB$OvP2f{~`{1W88a~g`qses9`*z%Ad2<&`TjR$a@Tx;< zMw{T2vyZpB+T$%n@*Lg^9MpwdfS5zeV*PyP!ofwc|EAb~@wp|B_=@n7Fwegn_<}>; z9ncQ{*Tnv>bh2V^Mz*?s=NsU}dxC=`A?ilB2z*;`RIaDJ{3?Ngcfn}WH9iAha#lnJ z9)Paz42-9NNw-v}35+Lc44iXWff2O=Bf3gJKqjsR#5}yr)!YV?@Rba1gsjXusDG89 z;F8?`H7Gbt?g7JbZe=pu1I=;1CMw8w#WWyzpdcXZc!=B+50U_`fFZwxZ%xq|f!fZ@?wSgQk#5WwdlEWqMPoYzaR0kh?wE zYq*s>;_xuN_IHzfZXA 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()