在node.js环境中简单实现一下令牌桶算法
今天突然想到了这个问题，公共接口的Api是怎么做接口限流的，然后思考了之后，刚开始想到了一种 固定窗口计数器算法，但是这个方法有个不好的点在于，时间边缘处不太友好，举个例子来说1秒钟允许5次请求的话，那么我地 990毫秒并发了五个请求，然后第1010秒又可以发五个请求，这样的话在20ms的时间里可以并发十个请求，虽然是每秒允许五个，但这种短时间内最大并发十个可能
并不是太好的选择。
于是我问了Ai，Ai提到了令牌桶算法，虽然名字很抽象，但是概念并不是很难理解，简单来说就是
根据pass time 与最大阈值来阶段性的生成可允许访问的数量，用代码抽象出来的话就是这样的
export class Bucket {
private capacity: number; // 桶容量
public current: number; // 当前水量
private interval: number; // 补水速率
private lastTime: number; // 补水时间
constructor(capacity: number, interval: number) {
this.capacity = capacity;
this.current = capacity;
this.interval = interval;
this.lastTime = Date.now();
}
｝这里最关键的是理解capacity和interval字段的含义，capacity就是充当了阈值，而interval代表了“流速”（生成令牌的速度），比如一秒允许五次，那么0.8秒就生成了 5*0.8 = 4 个令牌，而current代表当前的令牌量（水桶中水的体积）,
那么用户每次在接口的访问，就会消耗一个令牌，然后在消耗令牌之前会先进行补充，因为每次请求之间理论上是有时间间隔的，哪怕很小，所以，需要先补充这段时间间隔产生的令牌量，有可能是零点几个，取决于接口并发访问的时间
所以这里来把方法定义一下
refill() {
const now = Date.now();
const time = Number(((now - this.lastTime) / 1000).toFixed(1));
const water = time * this.interval; // 需要补水量
console.log("water:", water);
this.current = Math.min(this.current + water, this.capacity);
console.log("current:", this.current);
this.lastTime = now;
}
/**
* @description 消耗
*/
consume() {
this.refill();
if (this.current >= 1) {
this.current--;
return true;
}
return false;
}
这里consume方法返回值，是因为我要根据用户来区分桶，不同的用户有不同的桶，所以我还要再抽象一个用户控制的类，这里采用内存存储，在生产环境中，可以使用db数据库或者redis缓存数据库来存储，当然为了方便我这里采用内存来达到最小可行性的目的
import { NextFunction, Response, Request } from "express";
import { Bucket } from "./Bucket";
class UserCircumScribe {
static UserData: Map<string, Bucket> = new Map();
middleware(req: Request, res: Response, next: NextFunction) {
const userId = req.header("user-id") as string;
if (!userId)
return res.status(400).json({
code: 400,
message: "user-id is required",
});
if (!UserCircumScribe.UserData.has(userId)) {
UserCircumScribe.UserData.set(userId, new Bucket(1, 0.1));
}
const userBucket = UserCircumScribe.UserData.get(userId);
if (userBucket) {
const pass = userBucket.consume();
if (pass) {
return next();
}
res.status(429).json({
code: 429,
message: "请求次数超出限制",
});
}
}
}
export const circumscribe = new UserCircumScribe();因为我使用的是express.js框架，那么我就可以利用express.js的中间件机制来加载到某个接口或者全部接口上
// app.ts
import express, { Request, Response } from "express";
import { circumscribe } from "./lib/UserCircumscribe";
const app = express();
app.get("/", (req: express.Request, res: express.Response) => {
res.send("Hello World!");
});
app.get(
"/restriction",
circumscribe.middleware,
(_: Request, res: Response) => {
res.status(200).json({
message: "一些数据",
});
}
);
app.listen(3000);OK，我这里为了方便测试，设置的桶容量默认是1，流速是0.1，意思是十秒钟才生成一个令牌，看看效果
这样的话，一个令牌桶算法就在node.js中实现完成了，不过express生态中有提供的现成的包可以用：
express-rate-limit - npm
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！
页:
[1]