| from __future__ import annotations |
|
|
| import threading |
| import time |
| from dataclasses import dataclass |
|
|
|
|
| @dataclass |
| class _Bucket: |
| capacity: float |
| tokens: float |
| refill_per_sec: float |
| last_refill: float |
|
|
|
|
| class TokenBucket: |
| """In-memory token-bucket rate limiter keyed by string. |
| |
| Each key gets its own bucket lazily. Restart-safe by design: limits reset |
| when the Space restarts (acceptable per the design's in-memory state model). |
| """ |
|
|
| def __init__(self, capacity: int, refill_per_minute: int): |
| if capacity <= 0 or refill_per_minute <= 0: |
| raise ValueError("capacity and refill_per_minute must be positive") |
| self._capacity = float(capacity) |
| self._refill_per_sec = refill_per_minute / 60.0 |
| self._buckets: dict[str, _Bucket] = {} |
| self._lock = threading.Lock() |
|
|
| def _get(self, key: str, now: float) -> _Bucket: |
| b = self._buckets.get(key) |
| if b is None: |
| b = _Bucket( |
| capacity=self._capacity, |
| tokens=self._capacity, |
| refill_per_sec=self._refill_per_sec, |
| last_refill=now, |
| ) |
| self._buckets[key] = b |
| return b |
|
|
| def try_consume(self, key: str, n: int = 1) -> tuple[bool, int]: |
| """Attempt to consume n tokens. Returns (allowed, retry_after_seconds).""" |
| now = time.monotonic() |
| with self._lock: |
| b = self._get(key, now) |
| elapsed = now - b.last_refill |
| b.tokens = min(b.capacity, b.tokens + elapsed * b.refill_per_sec) |
| b.last_refill = now |
| if b.tokens >= n: |
| b.tokens -= n |
| return True, 0 |
| deficit = n - b.tokens |
| retry_after = max(1, int(deficit / b.refill_per_sec) + 1) |
| return False, retry_after |
|
|
| def refund(self, key: str, n: int = 1) -> None: |
| """Return n tokens to key's bucket, capped at capacity. Used when a |
| sibling bucket in a CompoundLimiter rejected after this one consumed.""" |
| with self._lock: |
| b = self._buckets.get(key) |
| if b is None: |
| return |
| b.tokens = min(b.capacity, b.tokens + n) |
|
|
|
|
| class CompoundLimiter: |
| """Apply multiple TokenBuckets to the same key; the strictest wins. |
| |
| A rejection must not burn tokens in the buckets that did allow — otherwise |
| a client hammering past one limit would also drain its allowance under the |
| others and be throttled harder than configured — so on rejection the |
| already-consumed buckets are refunded. |
| """ |
|
|
| def __init__(self, *buckets: TokenBucket): |
| self._buckets = buckets |
|
|
| def try_consume(self, key: str, n: int = 1) -> tuple[bool, int]: |
| consumed: list[TokenBucket] = [] |
| worst_retry = 0 |
| for b in self._buckets: |
| allowed, retry = b.try_consume(key, n) |
| if allowed: |
| consumed.append(b) |
| else: |
| worst_retry = max(worst_retry, retry) |
| if worst_retry: |
| for b in consumed: |
| b.refund(key, n) |
| return False, worst_retry |
| return True, 0 |
|
|
