Files
xiaobai-review/xiaobai-datahub/datahub/governance/ratelimit.py
T
3498dd7a4b feat(HEL-382): 搭建 datahub 底座和盘后正式数据链路
新增独立 xiaobai-datahub 服务(SQLite WAL、Tushare 盘后发布、/v1 契约和管理后台),不改现站页面与数据链路。

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: multica-agent <github@multica.ai>
2026-09-02 12:05:26 +08:00

37 lines
1.3 KiB
Python

from __future__ import annotations
import threading
import time
class TokenBucket:
def __init__(self, rate_per_minute: float, capacity: float | None = None, clock=time.monotonic) -> None:
self.rate_per_second = max(0.001, rate_per_minute / 60.0)
self.capacity = float(capacity if capacity is not None else rate_per_minute)
self._tokens = self.capacity
self._updated = clock()
self._clock = clock
self._lock = threading.Lock()
def acquire(self, tokens: float = 1.0, block: bool = True) -> bool:
while True:
with self._lock:
now = self._clock()
elapsed = max(0.0, now - self._updated)
self._tokens = min(self.capacity, self._tokens + elapsed * self.rate_per_second)
self._updated = now
if self._tokens >= tokens:
self._tokens -= tokens
return True
wait = (tokens - self._tokens) / self.rate_per_second
if not block:
return False
time.sleep(min(wait, 0.05))
@property
def remaining(self) -> float:
with self._lock:
now = self._clock()
elapsed = max(0.0, now - self._updated)
return min(self.capacity, self._tokens + elapsed * self.rate_per_second)