File size: 3,178 Bytes
30ea0e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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
|