用Python标准库从零实现令牌桶限流器

限流是后端避不开的话题。Python标准库里的threading和time就足够写出一个正确的令牌桶限流器。

一、令牌桶 vs 漏桶

漏桶输出速率恒定，令牌桶允许突发。对面向用户的接口，突发容忍更合理。

二、核心实现

惰性计算：记录上次补充时间，本次请求时根据时间差一次性补足。

import threading
import time

class TokenBucket:
    def __init__(self, rate, capacity):
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.last = time.monotonic()
        self.lock = threading.Lock()

    def acquire(self, need=1):
        with self.lock:
            now = time.monotonic()
            elapsed = now - self.last
            self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
            self.last = now
            if self.tokens >= need:
                self.tokens -= need
                return True
            return False

三、装饰器封装

def rate_limit(rate, capacity, on_reject=None):
    bucket = TokenBucket(rate, capacity)
    def deco(fn):
        def wrapper(*args, **kwargs):
            if bucket.acquire():
                return fn(*args, **kwargs)
            if on_reject is not None:
                return on_reject(*args, **kwargs)
            raise RuntimeError('rate limited')
        return wrapper
    return deco

四、并发细节

1. 用time.monotonic()不用time.time()。2. 锁要覆盖读-判-写整体。3. 不要sleep等待。

五、小结

整个实现不到50行，无第三方依赖，行为可预测。