API限流算法引擎·专业版

输入 / 输出

输入：限流参数（容量/速率或窗口）、请求时间；输出：allow() 是否放行、retry_after() 等待秒数、current_count() 窗口计数。

使用示例

from rate_limiter import TokenBucket
import time
tb = TokenBucket(capacity=100, refill_rate=10, clock=time.time)
if tb.allow():
    ...  # 处理请求
else:
    ...  # 返回 429，Retry-After: tb.retry_after()

引擎代码：rate_limiter.py

# -*- coding: utf-8 -*-
"""API 限流算法引擎：令牌桶 / 漏桶 / 滑动窗口。确定性算法，时间戳由调用方传入，便于测试与复现。"""


class TokenBucket:
    """令牌桶：允许一定突发（桶容量内），按固定速率补充令牌。

    capacity: 桶容量（最大突发请求数）
    refill_rate: 每秒补充令牌数
    """

    def __init__(self, capacity, refill_rate, clock=lambda: 0.0):
        if capacity <= 0 or refill_rate <= 0:
            raise ValueError("capacity 与 refill_rate 必须为正数")
        self.capacity = float(capacity)
        self.refill_rate = float(refill_rate)
        self._clock = clock
        self.tokens = float(capacity)
        self.last = clock()

    def allow(self, tokens=1):
        now = self._clock()
        if now < self.last:
            raise ValueError("时钟不能回退")
        self.tokens = min(self.capacity, self.tokens + (now - self.last) * self.refill_rate)
        self.last = now
        if self.tokens >= tokens:
            self.tokens -= tokens
            return True
        return False

    def retry_after(self, tokens=1):
        """距下一次放行还需等待的秒数（当前令牌不足时）。"""
        need = tokens - self.tokens
        return max(0.0, need / self.refill_rate)


class LeakyBucket:
    """漏桶：请求以固定速率漏出，超出容量直接拒绝，输出速率恒定（适合平滑流量）。"""

    def __init__(self, capacity, leak_rate, clock=lambda: 0.0):
        if capacity <= 0 or leak_rate <= 0:
            raise ValueError("capacity 与 leak_rate 必须为正数")
        self.capacity = float(capacity)
        self.leak_rate = float(leak_rate)
        self._clock = clock
        self.water = 0.0
        self.last = clock()

    def allow(self, amount=1):
        now = self._clock()
        if now < self.last:
            raise ValueError("时钟不能回退")
        self.water = max(0.0, self.water - (now - self.last) * self.leak_rate)
        self.last = now
        if self.water + amount <= self.capacity:
            self.water += amount
            return True
        return False


class SlidingWindowLog:
    """滑动窗口日志：记录窗口内每个请求的时间戳，统计窗口内请求数，边界精确（无固定窗口的2倍突发问题）。"""

    def __init__(self, max_requests, window_seconds, clock=lambda: 0.0):
        if max_requests <= 0 or window_seconds <= 0:
            raise ValueError("max_requests 与 window_seconds 必须为正数")
        self.max_requests = int(max_requests)
        self.window = float(window_seconds)
        self._clock = clock
        self.log = []

    def allow(self):
        now = self._clock()
        boundary = now - self.window
        self.log = [t for t in self.log if t > boundary]
        if len(self.log) < self.max_requests:
            self.log.append(now)
            return True
        return False

    def current_count(self):
        now = self._clock()
        return len([t for t in self.log if t > now - self.window])

测试（9 正例 + 3 边界（test_rate_limiter.py 实跑全过）：满桶突发、时间推进补令牌、漏桶恒定漏出、滑动窗口精确计数与无2倍突发；非法参数、时钟回退、零配额边界均正确报错。）

硬性规则

• 确定性输出：相同输入必得相同结果，时间/随机量通过参数注入；
• 非法输入（负数、越界、格式错误）一律抛 ValueError，不静默吞掉；
• 本引擎为算法/计算参考，集成到业务时请结合真实参数与压测，不替代专业判断。