简单地实现令牌桶限速

实现

1. 按一定的速率往令牌桶添加令牌

2. 只有获取令牌才能处理任务。

3. 令牌桶满后，新增的令牌丢弃。

4. 消费令牌后，发现令牌桶没有令牌，及时把消费的令牌还给令牌桶。

优点

滑动窗口：允许在最小窗口的边界，超过限定的速率。例如：按最小窗口为分钟，每分钟最大消费60次，在01分59秒和02分0秒这2秒的时间内，允许消费120次，其他时间不消费。

漏桶算法：当漏桶满以后，只能按一定速率消费。无法过载处理，缺乏效率。

令牌算法：允许过载处理，可以一下子把令牌桶的令牌消费，由于按一定速率添加令牌，令牌消费按照一定速率限速。消费令牌后，发现令牌桶没有令牌，及时补充令牌，防止添加令牌过慢，系统相对空闲的情况。

interface Options {
    interval: number;
    max: number;
}

function wait(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
}
 
// 简单地实现令牌桶限速
export class RateLimit {
    private _interval = 60000;
    private _max = 100;
    private _buckets = new Map<string, { count: number, lastTime: number }>();
    constructor(options: Options) {
        this.interval = options.interval;
        this.max = options.max;
    }
 
    get max() {
        return this._max
    }
    set max(val: number) {
        if (val > 0) {
            this._max = Math.floor(val);
        }
    }
 
    get interval() {
        return this._interval;
    }
    set interval(val: number) {
        if (val >= 1000) {
            this._interval = Math.floor(val);
        }
    }

    get tokenRate() {
        return this._max / this._interval;
    }

    getBucket(key: string) {
        let bucket = this._buckets.get(key);
        const now = Date.now();
        if (!bucket) {
            bucket = { count: 0, lastTime: now };
            this._buckets.set(key, bucket);
        }
        if (now < bucket.lastTime) {
            bucket.lastTime = now;
        }
        bucket.count += (now - bucket.lastTime) * this.tokenRate;
        bucket.lastTime = now;
        if (bucket.count > this._max) {
            bucket.count = this._max;
        }
        return bucket;
    }
 
    async limit(key: string, count = 1) {
        if (count <= 0) {
            throw new Error(`${count} <= 0`);
        }
        if (count > this._max) {
            throw new Error('Count is greater than the maximum value')
        }
        const bucket = this.getBucket(key);
        if (bucket.count - count >= 0) {
            bucket.count -= count;
        } else {
            await wait((count - bucket.count) / this.tokenRate);
            return await this.limit(key, count);
        }
        return bucket.count;
    }

    tryLimit(key: string, count = 1) {
        if (count <= 0 || count > this._max) {
            return false;
        }
        const bucket = this.getBucket(key)
        if (count > bucket.count) {
            return false;
        }
        bucket.count -= count;
        return true;
    }
 
    consume(key: string, count: number) {
        if (count <= 0) {
            return;
        }
        const bucket = this._buckets.get(key);
        if (!bucket) {
            return;
        }
        // 令牌桶为空，补充令牌
        if (bucket.count == 0) {
            bucket.count += count;
        }
    }
}