用 Python 标准库从零实现一个生产可用的令牌桶限流器

限流是后端避不开的话题。Python 生态里限流库不少，但大多数场景下你其实不需要引入它们——标准库里的 threading 和 time 就足够写出一个正确、轻量、好理解的令牌桶限流器。这篇笔记把整个实现拆开讲清楚，重点放在几个并发上容易写错的细节。

一、令牌桶 vs 漏桶：别选错

两种算法常被混为一谈，但行为差别很大。漏桶（Leaky Bucket）以固定速率漏水，请求像倒水一样进来，桶满了直接拒绝——它的输出速率永远恒定，不管你瞬间倒多少水。令牌桶（Token Bucket）则是以固定速率往桶里放令牌，请求来了先拿令牌，拿得到就放行，桶里的令牌可以攒到一个上限。

关键差别：令牌桶允许突发。如果桶容量是 10、填充速率 5/秒，那么空闲一段时间后桶攒满 10 个令牌，下一秒可以瞬间放过 10 个请求，之后才回落到 5/秒。对大多数面向用户的接口，这种突发容忍反而更合理——用户刷新页面会并发打几个接口，漏桶会误杀，令牌桶不会。

二、核心实现

实现的关键是不真的起一个线程去定时放令牌，而是惰性计算：记录上次补充时间，本次请求时根据时间差一次性补足。这样实现零额外线程，也避免了定时器的漂移问题。

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

不到 20 行。几个要点：用 time.monotonic() 而不是 time.time()，因为限流算的是间隔，墙上时钟会被 NTP 调整甚至回拨，monotonic 保证单调递增；补充量用 min 截断到 capacity，否则空闲一天能攒出几百万令牌，瞬间放行就把下游打挂；锁只保护这一小段计算，临界区极短。

三、装饰器封装

直接调 acquire 不够优雅，包成装饰器更顺手。这里一个常见需求是"被限流时怎么办"——直接抛异常、返回降级值、还是 sleep 等待。把策略做成可注入的回调最灵活：

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

@rate_limit(rate=5, capacity=10, on_reject=lambda *a, **k: {'msg': 'too busy'})
def handle(req):
    ...

注意桶是装饰器闭包里的单例，所有对同一函数的调用共享一个桶，这正是我们想要的。如果给每个调用单独建桶，限流就形同虚设了。

四、几个容易写错的并发细节

1. 别用 time.time()。前面提过，墙上时钟回拨会让 elapsed 变成负数，tokens 越算越少，限流器直接卡死。monotonic 是为这种场景生的。

2. 锁的粒度。有人为了"减少锁竞争"把 tokens 的读取放在锁外、只把写回放锁内，这是经典的 check-then-act 竞态：两个线程同时读到 tokens=1，都判断够用，都进临界区减一，tokens 变成 -1。限流器必须把"读-判-写"作为一个原子操作整体放进锁里。

3. acquire 的 need 参数。有些接口一个请求消耗多个令牌（比如批量接口按批量大小计费），把 need 暴露出来比写多个限流器更干净。但要注意：当 need > capacity 时永远拿不到令牌，构造时应断言 need <= capacity，否则是个静默死锁。

4. 不要 sleep 等待。有的实现 acquire 失败会 time.sleep 到下一个令牌可用再重试。在 Web 上下文里这很危险——线程被占着 sleep，并发一高线程池就耗光了。宁可快速失败让上层决定重试或降级。

五、验证与小结

验证最直观的办法是起 100 个线程同时抢一个 rate=10/capacity=10 的桶，统计 1 秒内放行数——应该先瞬间放 10 个，之后大约每秒 10 个。实测下来和理论吻合，抖动在 1 个以内。

整个实现不到 50 行，没有第三方依赖，行为可预测。当然它只适合单机限流；分布式限流得把 tokens 状态搬到 Redis 之类的地方，用 SET + Lua 脚本保证原子性，那是另一个话题了。但即便上分布式，单机这层令牌桶作为本地预过滤仍然值得保留——能挡掉大部分无效请求，不让它们打到 Redis 上。