限流

token_bucket.py

import time


class TokenBucket:
    """Allow bursts of up to `capacity` requests, refilled at `rate` tokens per second."""

    def __init__(self, rate: float, capacity: float) -> None:
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.updated = time.monotonic()

    def allow(self, cost: float = 1.0) -> bool:
        now = time.monotonic()
        self.tokens = min(self.capacity, self.tokens + (now - self.updated) * self.rate)
        self.updated = now
        if self.tokens >= cost:
            self.tokens -= cost
            return True
        return False


if __name__ == "__main__":
    bucket = TokenBucket(rate=2, capacity=5)   # 2 requests per second, bursts of 5
    print([bucket.allow() for _ in range(8)])  # 5 x True, then 3 x False
    time.sleep(1.1)                            # about 2.2 tokens come back
    print(bucket.allow(), bucket.allow(), bucket.allow())  # True True False

python token_bucket.py