限流算法学习：漏桶 & 令牌桶算法

引言

本节主要学习下两种常用的单机限流思想，分别是漏桶算法和令牌桶算法。此外，还将给出使用 Python 及 Go 语言实现，便于加深理解。当然，现实中肯定不能直接用下面的代码。实际应用时，我们不大可能在单机执行限流，下面的实现也并非线程或 goroutine 安全的。

实际限流可以考虑在 Nginx 层对请求限流，或者如果真的要自己在业务方实现一套限流策略的话，可以考虑基于 Redis 实现分布式限流策略。并且在实际应用中，可能还会基于不同的维度进行限流，如用户 id，请求 IP 等，实际应用需要考虑的东西更多。

漏桶算法

可以把请求当作水流，水流全部进入有限大小的水缸，同时水缸会按照固定的速率漏水。当水流湍急，水缸漏水太慢的话，会得知水缸积水，直到溢出（此时拒绝服务）。

特点

1. 实现起来很简单，并且能够以比较恒定的速率服务请求
2. 缺点是无法应对突发流量，很容易导致溢出

实现

Python

class LeakyBucketRateLimiter(object):
    def __init__(self, capacity=10, leak_rate=1):
        """
        初始化漏桶
        :param capacity: 桶容量
        :param leak_rate: 恒定的消费速度（Reqs/秒）
        """
        self._capacity = float(capacity)
        self._leak_rate = float(leak_rate)
        self._water_level = 0.0
        # 上次漏水的时间
        self._last_time = time.time()

    def acquire(self, level=1):
        # 执行漏水
        now = time.time()
        delta = self._water_level - self._leak_rate * (now - self._last_time)
        self._water_level = min(0.0, delta)
        self._last_time = now
        # 尝试加水，并看水桶是否满了
        if level + self._water_level > self._capacity:
            raise RateLimitExceeded()
        self._water_level += level

Go

type LeakyBucketRateLimiter struct {
    capacity int
    currentLevel int
    leakRate int // consume how many requests per sec
    lastLeakedAt time.Time
}

func NewLeakyBucketRateLimitter(capacity, leakRate int) *LeakyBucketRateLimiter {
    return &LeakyBucketRateLimiter{
        capacity: capacity,
        currentLevel: 0,
        leakRate: leakRate,
        lastLeakedAt: time.Now(),
    }
}

func (r *LeakyBucketRateLimiter) Acquire(n int) error {
    now := time.Now()
    // leak water
    currentLevel := r.currentLevel - r.leakRate*int(now.Sub(r.lastLeakedAt).Seconds())
    r.currentLevel = max(currentLevel, 0)
    r.lastLeakedAt = now
    // try to add water, test bucket is full or not.
    currentLevel = n + r.currentLevel
    if currentLevel > r.capacity {
        return errRateLimitExceeds
    }
    r.currentLevel = currentLevel
    return nil
}

令牌桶算法

同样想象我们有一个桶，专门存放令牌，会以恒定的速率生成令牌，并将其放入桶中。每当有请求过来时，需要先从桶中取到一个或多个令牌，如果获取成功，则为请求提供服务，否则拒绝服务。

特点

1. 实现同样是很简单
2. 可以应对突发流量，面对瞬间大流量，可以在短时间内获得大量令牌，且生产令牌毫不费力
3. 可以做流量整形

实现

Python

class TokenBucketRateLimiter(object):
    def __init__(self, capacity=1, fill_rate=1):
        """
        初始化令牌桶限流器
        :param capacity: 令牌桶容量
        :param fill_rate: 放入令牌的速度（Reqs/秒）
        """
        self._capacity = float(capacity)
        self._rate = float(fill_rate)
        self._bucket_tokens = float(capacity)
        # 上次添加令牌的时间
        self._last_time = int(time.time())

    def acquire(self, tokens=1):
        # 发放令牌
        if self._bucket_tokens < self._capacity:
            now = time.time()
            delta = (now - self._last_time) * self._rate
            self._last_time = now
            self._bucket_tokens = min(self._capacity, self._bucket_tokens + delta)

        if tokens > self._bucket_tokens:
            # 无法获取令牌了，数量不够
            raise RateLimitExceeded()
        self._bucket_tokens -= tokens

Go

type TokenBucketRateLimiter struct {
    capacity int
    tokens int
    putRate int // put how many tokens per sec
    lastPutAt time.Time
}

func NewTokenBucketRateLimiter(capacity, fillRate int) *TokenBucketRateLimiter {
    return &TokenBucketRateLimiter{
        capacity: capacity,
        tokens: 0,
        putRate: fillRate,
        lastPutAt: time.Now(),
    }
}

func (r *TokenBucketRateLimiter) Acquire(n int) error {
    if r.tokens < r.capacity {
        // put tokens in the bucket
        now := time.Now()
        howMany := r.putRate * int(now.Sub(r.lastPutAt).Seconds())
        r.tokens = min(r.capacity, howMany+r.tokens)
        r.lastPutAt = now
    }
    // check if we have enough tokens
    if r.tokens < n {
        return errRateLimitExceeds
    }
    // release tokens
    r.tokens -= n
    return nil
}

参考

• 接口限流算法介绍