| import time | |
| from collections import defaultdict | |
| from typing import Optional | |
| class TokenBucket: | |
| """令牌桶限流器（每个 key 独立）""" | |
| def __init__(self, rate: float, burst: int): | |
| self.rate = rate # 每秒令牌数 | |
| self.burst = burst # 最大令牌数 | |
| self._tokens = defaultdict(float) | |
| self._last = defaultdict(float) | |
| def allow(self, key: str) -> bool: | |
| """返回 True 表示允许请求，False 表示限流""" | |
| now = time.time() | |
| # 补充令牌 | |
| last = self._last.get(key, now) | |
| elapsed = now - last | |
| self._last[key] = now | |
| self._tokens[key] = min(self.burst, self._tokens.get(key, self.burst) + elapsed * self.rate) | |
| if self._tokens[key] >= 1.0: | |
| self._tokens[key] -= 1.0 | |
| return True | |
| return False | |
| class RateLimiter: | |
| """多 key 限流器，qps<=0 时放行""" | |
| def __init__(self, qps: int = 0, burst: int = 0): | |
| self.qps = qps | |
| self.burst = burst if burst > 0 else qps | |
| self._bucket = TokenBucket(rate=qps, burst=self.burst) if qps > 0 else None | |
| def allow(self, key: str) -> bool: | |
| if self._bucket is None: | |
| return True | |
| return self._bucket.allow(key) | |
| def active_buckets(self) -> int: | |
| """返回当前活跃的 key 数量（近似）""" | |
| if self._bucket is None: | |
| return 0 | |
| return len(self._bucket._tokens) | |