基于python实现限流器-令牌桶算法

4 阅读4分钟

令牌桶(Token Bucket)算法

算法原理

核心思想:通过控制令牌的生成和消耗来限制系统的流量

  1. 令牌桶的结构:令牌桶是一个存储令牌的容器,它有一个固定的容量和一个固定的令牌生成速率。令牌桶通常由两个部分组成:令牌存储区和令牌生成器。令牌存储区用于存储生成的令牌,而令牌生成器则负责按照一定的速率生成令牌,并将生成的令牌放入令牌存储区中。
  2. 令牌的生成:令牌生成器按照固定的速率生成令牌,并将生成的令牌放入令牌存储区中。令牌生成的速率通常由系统管理员根据系统的实际需求进行设置。
  3. 令牌的消耗:当一个请求到达系统时,系统会从令牌存储区中取出一个令牌。如果令牌存储区中有足够的令牌,那么请求可以被立即处理;如果令牌存储区中没有足够的令牌,那么请求会被拒绝或者等待,直到令牌存储区中有足够的令牌为止。

令牌的生成方式

  1. 固定速率:令牌生成器按照固定的速率生成令牌。例如,如果系统需要限制每秒的请求数量为 100 个,那么令牌生成器可以设置为每秒生成 100 个令牌。
  2. 动态速率:令牌生成器可以选择随机生成令牌的速率。这种方式可以更加灵活地控制令牌的生成速率,但也可能导致请求被拒绝的概率增加。

令牌的消耗方式

  1. 一次性消耗:一次性消耗是最简单的令牌消耗方式。在这种方式下,当一个请求到达系统时,系统会从令牌存储区中取出一个令牌
  2. 动态消耗:分批消耗是一种更加灵活的令牌消耗方式。在这种方式下,当一个请求到达系统时,系统会从令牌存储区中取出多个令牌。比如get请求消耗一个令牌,post请求消耗两个令牌

算法特点

  1. 平滑限流:滑动窗口算法通过在一个固定时间窗口内逐步滑动子窗口,从而精细地控制请求的流量。相比固定窗口算法,滑动窗口算法能有效避免临界点突然放开导致的流量激增。
  2. 高适应性:滑动窗口算法适用于请求量变化较大的场景,因为它能够动态调整每个时间段的请求限制,保证整个窗口的请求数平稳。

存在问题

  1. 存储开销较大:滑动窗口算法通常需要记录多个子窗口的请求计数以计算滑动窗口内的总请求数。这会导致内存开销上升,尤其是在子窗口数量较多的情况下。
  2. 吞吐量不足:在频繁滑动的过程中,部分请求可能会因为算法过于严格的控制而被拒绝,导致流量利用不充分,尤其在流量突发的情况下可能出现“漏接”请求的情况。
  3. 边界效应:滑动窗口算法虽然平滑了请求分布,但在窗口边界处,仍然会存在一些无法精确控制的流量跳跃。例如,当滑动窗口正好跨越一个请求高峰区时,可能会导致短暂的流量堆积。

算法实现

import time

class TokenBucketRateLimiter:
    def __init__(self, capacity, refill_rate):
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.tokens = capacity
        self.last_refill = time.time()

    def allow_request(self):
        current_time = time.time()
        time_passed = current_time - self.last_refill
        refill_amount = time_passed * self.refill_rate
        self.tokens = min(self.capacity, self.tokens + refill_amount)
        self.last_refill = current_time
        if self.tokens >= 1:
            self.tokens -= 1
            return True
        return False

# 测试用例
def test_token_bucket_rate_limiter():
    limiter = TokenBucketRateLimiter(capacity=5, refill_rate=1)

    # 正常限流:请求速度与令牌生成速度相匹配
    
    for _ in range(5):
        print(limiter.allow_request())  # Expected: True
    time.sleep(1)
    print(limiter.allow_request())  # Expected: True
    
    time.sleep(5)
    # 突发请求超过令牌生成速率
    for _ in range(10):  # 超过容量的请求将被拒绝
        print(limiter.allow_request())  # Expected: True for first 5, then False

test_token_bucket_rate_limiter()