用令牌桶实现进程内接口限流：惰性补充、并发安全与时钟回拨处理

在单进程 Python 服务里，接口限流经常被写成“固定窗口计数”：每分钟清零一次。这种写法实现简单，但窗口边界附近可能出现双倍突发，且多线程同时读写计数时容易超额放行。令牌桶更适合表达“平均速率 + 突发容量”：桶里最多存 capacity 个令牌，系统按 rate 个/秒补充；请求需要先取令牌，取不到就拒绝，并告诉调用方大约等待多少秒。

本文实现一个进程内、线程安全的令牌桶，采用惰性补充：不启动后台线程，每次 acquire 时根据当前时间与上次补充时间计算新增令牌。默认使用 time.monotonic()，它不受系统墙钟调整影响；同时保留可注入时钟，便于测试时钟回拨。拒绝时返回 wait_seconds，调用方可据此返回 Retry-After 或延迟重试。

设计要点

• capacity 是突发上限，也是初始令牌数；rate 是补充速率，单位令牌/秒。
• 每次请求在锁内完成“补充 + 扣减/判断”，避免检查与扣减分离。
• 只有 now > last_refill 才补充，并把 last_refill 推进到 now；now <= last_refill 时不补充、不回退 last_refill，防止时钟回拨后重复补充。
• 拒绝等待时间：正常情况下为 need/rate；若 now < last_refill，还要加上时钟恢复到 last_refill 的时间，避免回拨时给出偏小的等待值。
• 单次请求 tokens 不能超过 capacity，否则永远无法满足，直接抛 ValueError。

最小实现

import threading
import time
from dataclasses import dataclass

@dataclass
class Decision:
    allowed: bool
    wait_seconds: float
    remaining: float

class TokenBucket:
    def __init__(self, rate: float, capacity: int, clock=time.monotonic):
        if rate <= 0:
            raise ValueError("rate must be > 0")
        if capacity < 1:
            raise ValueError("capacity must be >= 1")
        self._rate = float(rate)
        self._capacity = float(capacity)
        self._tokens = float(capacity)
        self._clock = clock
        self._last = clock()
        self._lock = threading.Lock()

    def _refill_locked(self, now: float) -> None:
        if now > self._last:
            delta = now - self._last
            self._tokens = min(self._capacity, self._tokens + delta * self._rate)
            self._last = now
        # now <= self._last：时钟回拨或同一时刻，不补充，也不回退 _last

    def acquire(self, tokens: float = 1.0) -> Decision:
        if tokens <= 0:
            raise ValueError("tokens must be > 0")
        if tokens > self._capacity:
            raise ValueError("tokens must not exceed capacity")
        with self._lock:
            now = self._clock()
            self._refill_locked(now)
            if self._tokens >= tokens:
                self._tokens -= tokens
                return Decision(True, 0.0, self._tokens)
            need = tokens - self._tokens
            wait = need / self._rate
            if now < self._last:
                wait += self._last - now
            return Decision(False, wait, self._tokens)

这里 Decision.remaining 是判断后的剩余令牌数；拒绝时它表示当前可用令牌。wait_seconds 是至少需要等待的估计值，并发竞争下其他线程可能先取走新令牌，实际等待可能更长。

运行方法与并发验证

将代码保存为 token_bucket.py。在文件末尾加入：

from concurrent.futures import ThreadPoolExecutor

if __name__ == "__main__":
    bucket = TokenBucket(rate=5, capacity=10)

    def call(i):
        d = bucket.acquire(1)
        return i, d.allowed, d.wait_seconds

    with ThreadPoolExecutor(max_workers=20) as ex:
        results = list(ex.map(call, range(50)))

    allowed = sum(1 for _, ok, _ in results if ok)
    rejected = len(results) - allowed
    waits = [round(w, 3) for _, ok, w in results if not ok]
    print("allowed:", allowed)
    print("rejected:", rejected)
    print("sample waits:", waits[:5])

运行 python token_bucket.py。输出数值会随机器调度和运行耗时变化，不应写死断言。验证思路是：初始桶容量为 10，50 个并发请求在很短时间内到达，成功数不会超过“10 + 运行耗时 × 5”的粗略上界；每个被拒绝请求的 wait_seconds 应为正数。若要更严格地验证突发上限，可先把 rate 设得很小（例如 0.001），让补充量在测试窗口内可忽略，再检查成功数不超过 capacity。注意这是测试构造，不代表线上性能。

时钟回拨处理

默认 time.monotonic() 不会因系统时间调整而回拨，生产环境应优先使用它。为了验证回拨分支，可以注入一个假时钟：

class FakeClock:
    def __init__(self, t=100.0):
        self.t = t

    def __call__(self):
        return self.t

if __name__ == "__main__":
    c = FakeClock()
    b = TokenBucket(rate=1, capacity=2, clock=c)
    print(b.acquire(2))   # 允许，令牌清零
    c.t -= 10             # 模拟墙钟回拨 10 秒
    d = b.acquire(1)
    print(d.allowed, round(d.wait_seconds, 3))  # 拒绝，等待约 11 秒
    c.t += 11             # 回到上次补充时间之后 1 秒
    print(b.acquire(1).allowed)  # 允许

回拨期间不补充令牌，且不把 _last 回退到更早时间；拒绝等待时间包含“时钟恢复到 _last”的时长，因此不会因为回拨而低估等待。若业务必须使用墙钟，建议在发现负间隔时记录告警并考虑切换单调时钟；不要直接 last = now，否则时钟恢复后可能重复补充。

边界与限制

• 进程内限流只约束当前进程；多进程、多实例部署需要集中式存储或网关限流，不能把本实现当作全局限流。
• 不保证请求公平性，锁的调度顺序由运行时决定；高并发下锁竞争可能成为瓶颈。
• 桶满后多余补充时间被丢弃，不会累积超过 capacity 的突发额度。
• 拒绝等待秒数是估计值，调用方应允许重试，并设置上限；并发下实际等待可能更长。
• 请求 tokens 大于 capacity 会直接报错；如需按权重限流，应保证单次权重不超过容量。
• 进程重启后令牌桶回到初始满桶状态；如需持久化状态，应另行设计。
• 默认单调时钟不会回拨；注入墙钟时回拨分支偏保守，极端回拨会造成较长等待。

小结

令牌桶的核心不是后台定时器，而是每次请求时用“当前时间 - 上次补充时间”惰性计算令牌。把补充、扣减和判断放进同一把锁，处理好负时间间隔和拒绝等待时间，就能得到一个可读、可测试的进程内限流器。