后端接口限流设计:令牌桶、漏桶与滑动窗口的代码实现

后端接口限流设计:令牌桶、漏桶与滑动窗口的代码实现

后端接口限流设计:令牌桶、漏桶与滑动窗口的代码实现

一、接口限流不是"加个计数器就行"

首先想象的场景很朴素:统计这个 IP 过去 1 分钟内请求了几次,超过 100 次就拒绝。这看起来只需要一个计数器。但实际场景比这个复杂得多:

  • 如果流量突然涌入,你是全部拒绝还是允许一些通过?(令牌桶解决的问题)
  • 如果流量持续均匀,你是平滑通过还是允许突发?(漏桶解决的问题)
  • 如何避免计数器边界附近的"双倍流量"?(滑动窗口解决的问题)

二、四种限流算法的本质区别

flowchart TD A[限流算法] --> B[固定窗口] A --> C[滑动窗口] A --> D[漏桶] A --> E[令牌桶] B --> B1["简单,但有边界双倍流量问题"] C --> C1["精确平滑,实现稍复杂"] D --> D1["强制平滑,不可突发"] E --> E1["允许突发,平滑平均"]
算法突发允许实现复杂度适用场景
固定窗口边界处允许简单场景,精度要求不高
滑动窗口不允许需要精确控制
漏桶不允许流量整型
令牌桶允许(桶容量内)允许合理突发

三、实现

import time import threading from collections import deque from dataclasses import dataclass @dataclass class RateLimitResult: """限流结果""" allowed: bool remaining: int # 剩余配额 reset_at: float # 重置时间(Unix 时间戳) class TokenBucket: """令牌桶限流器 核心思想: - 以固定速率向桶中添加令牌 - 桶有容量上限(允许的突发量) - 请求需要消费令牌,无令牌时拒绝 适用场景: - 允许合理突发的 API 限流 - 如:每秒 100 个请求,但允许瞬时 200 个(桶容量=200) 时间复杂度:O(1) 空间复杂度:O(1) """ def __init__(self, rate: float, capacity: int): """ Args: rate: 令牌生成速率(每秒生成多少个令牌) capacity: 桶的最大容量(允许的最大突发请求数) """ self.rate = rate self.capacity = capacity self._tokens = float(capacity) # 当前令牌数 self._last_refill = time.monotonic() self._lock = threading.Lock() def allow_request(self, tokens: int = 1) -> RateLimitResult: """尝试消费 tokens 个令牌""" with self._lock: self._refill() if self._tokens >= tokens: self._tokens -= tokens return RateLimitResult( allowed=True, remaining=int(self._tokens), reset_at=time.time() + self._time_to_refill(tokens), ) else: return RateLimitResult( allowed=False, remaining=0, reset_at=time.time() + self._time_to_refill(tokens - self._tokens), ) def _refill(self) -> None: """补充令牌""" now = time.monotonic() elapsed = now - self._last_refill # 按速率补充令牌,不超过容量上限 self._tokens = min(self.capacity, self._tokens + elapsed * self.rate) self._last_refill = now def _time_to_refill(self, needed: float) -> float: """计算需要等待多少秒才能获得 needed 个令牌""" return needed / self.rate if self.rate > 0 else float("inf") class SlidingWindowLog: """滑动窗口日志限流器 核心思想: - 维护一个时间戳队列 - 每次请求记录时间戳 - 判断时清理窗口外的旧记录,统计窗口内请求数 优势:精确,无边界双倍流量问题 劣势:需要存储窗口内所有请求的时间戳(内存开销大) 时间复杂度:O(1) 摊销(清理 + 入队) 空间复杂度:O(窗口内最大请求数) """ def __init__(self, max_requests: int, window_seconds: float): self.max_requests = max_requests self.window = window_seconds self._timestamps: deque[float] = deque() self._lock = threading.Lock() def allow_request(self) -> RateLimitResult: """判断当前请求是否被允许""" with self._lock: now = time.monotonic() # 清理窗口外的旧记录 window_start = now - self.window while self._timestamps and self._timestamps[0] <= window_start: self._timestamps.popleft() if len(self._timestamps) < self.max_requests: self._timestamps.append(now) return RateLimitResult( allowed=True, remaining=self.max_requests - len(self._timestamps), reset_at=time.time() + self.window, ) else: # 计算下次允许请求的时间 next_available = self._timestamps[0] + self.window wait = max(0, next_available - now) return RateLimitResult( allowed=False, remaining=0, reset_at=time.time() + wait, ) class SlidingWindowCounter: """滑动窗口计数器限流器(优化版) 与 SlidingWindowLog 的区别: - 不存储每个请求的时间戳,只统计上一窗口和当前窗口的计数 - 用加权平均估算滑动窗口内的请求数 - 内存开销 O(1),精度略低于日志版 适用场景:高并发下内存敏感的限流 """ def __init__(self, max_requests: int, window_seconds: float): self.max_requests = max_requests self.window = window_seconds self._prev_count = 0 # 上一个完整窗口的计数 self._curr_count = 0 # 当前窗口的计数 self._window_start = time.monotonic() self._lock = threading.Lock() def allow_request(self) -> RateLimitResult: """判断当前请求是否被允许""" with self._lock: now = time.monotonic() elapsed = now - self._window_start # 如果已经进入下一个窗口,滚动窗口 if elapsed >= self.window: self._prev_count = self._curr_count self._curr_count = 0 self._window_start = now elapsed = 0 # 加权计算:上一窗口的权重随已过时间递减 weight_prev = (self.window - elapsed) / self.window estimated = self._prev_count * weight_prev + self._curr_count if estimated < self.max_requests: self._curr_count += 1 return RateLimitResult( allowed=True, remaining=max(0, self.max_requests - int(estimated) - 1), reset_at=time.time() + (self.window - elapsed), ) else: wait = self.window - elapsed return RateLimitResult( allowed=False, remaining=0, reset_at=time.time() + wait, )

四、选型建议

场景推荐算法
简单的 API 限流(允许短时突发)令牌桶
第三方 API 调用频率控制令牌桶(设置桶容量为 1 禁止突发)
高精度防刷限流滑动窗口日志
高并发 + 内存敏感滑动窗口计数器
流量整型(消息队列消费控制)漏桶

五、总结

限流算法的选择依赖三个条件:是否允许突发、精度要求有多高、并发量有多大。固定窗口最简单但有边界问题,滑动窗口精确但内存开销大,令牌桶和漏桶用桶容量参数灵活控制突发行为。在大多数后端 API 限流场景中,令牌桶是综合最优解。