实现
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;
}
}
}