Token Bucket Rate Limiter
Implement a token bucket that refills at a steady rate and allows short bursts, a standard algorithm for shaping request traffic.
Code
import time
from dataclasses import dataclass, field
@dataclass
class TokenBucket:
capacity: float
refill_rate: float
tokens: float = field(init=False)
updated_at: float = field(init=False)
def __post_init__(self) -> None:
self.tokens = self.capacity
self.updated_at = time.monotonic()
def allow(self, cost: float = 1.0) -> bool:
now = time.monotonic()
elapsed = now - self.updated_at
self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
self.updated_at = now
if self.tokens >= cost:
self.tokens -= cost
return True
return FalseReport a problem with this exampleOpens a pre-filled issue on GitHub.