关于令牌桶和漏桶算法的碎碎念

Sep 25, 2023 · 2 min read · rate limit algorithm token bucket leaky ·

该文章根据 CC-BY-4.0 协议发表，转载请遵循该协议。
本文地址：https://fenying.net/post/2023/09/25/idea-about-rate-limit-algorithms/

关于漏桶算法和令牌桶算法的一些理解。

昨晚抽空看了一会令牌桶（Token Bucket）和漏桶（Leaky Bucket）两个算法。

想着令牌桶那么简单，每个实例只需要两个变量就够了，伪代码：

 1class TokenBucket {
 2
 3    private _lastRefillTime: number = Date.now();
 4    private _tokens: number;
 5
 6    constructor(
 7        private readonly _capacity: number,
 8        private readonly _refillInterval: number,
 9    ) {
10
11        this._tokens = _capacity;
12    }
13
14    private _tryRefill(): void {
15
16        const now = Date.now();
17
18        const elapsed = now - this._lastRefillTime;
19
20        const refillTokens = Math.floor(elapsed / this._refillInterval);
21
22        if (refillTokens > 0) {
23
24            this._tokens = Math.min(this._capacity, this._tokens + refillTokens);
25            this._lastRefillTime = now;
26        }
27    }
28
29    public consume(): boolean {
30
31        this._tryRefill();
32
33        if (!this._tokens) {
34
35            return false;
36        }
37
38        this._tokens--;
39
40        return true;
41    }
42}

而漏桶却那么复杂，还得弄个队列？

仔细想了下，貌似漏桶只需要关注什么时候漏出去，且漏桶的流出速率是固定的，也就是说——只需要关注当前桶里最后一个元素什么时候会漏出去。根据这个时间，一可以算出当前桶的大小，也可以算出下一个新元素什么时候漏出去，只需要让当前请求等到那个时候再执行即可。那每个实例只需要一个变量就够了，比令牌桶还简单……😓

 1async function wait(ms: number): Promise<void> {
 2
 3    return new Promise(resolve => setTimeout(resolve, ms));
 4}
 5
 6class LeakyBucket {
 7
 8    private _nextWaitUntil: number = Date.now();
 9
10    public constructor(
11        private readonly _capacity: number,
12        private readonly _leakInterval: number,
13    ) {}
14
15    public async leak(): Promise<void> {
16
17        const now = Date.now();
18        const elapsed = this._nextWaitUntil - now;
19
20        if (elapsed <= 0) { // 当前桶里的元素已经漏完了，也要等一个周期才能漏出去。
21
22            this._nextWaitUntil = now + this._leakInterval;
23            await wait(this._leakInterval);
24            return;
25        }
26
27        const bucketSize = Math.ceil(elapsed / this._leakInterval);  // 当前桶里的元素个数
28
29        if (bucketSize > this._capacity) {
30
31            // 桶满了，拒绝请求
32            throw new Error('Queue overflow');
33        }
34
35        this._nextWaitUntil = this._nextWaitUntil + this._leakInterval;
36        await wait(this._nextWaitUntil - now);
37    }
38}