在node.js环境中简单实现一下令牌桶算法

今天突然想到了这个问题，公共接口的Api是怎么做接口限流的，然后思考了之后，刚开始想到了一种 固定窗口计数器算法，但是这个方法有个不好的点在于，时间边缘处不太友好，举个例子来说
1秒钟允许5次请求的话，那么我地 990毫秒并发了五个请求，然后第1010秒又可以发五个请求，这样的话在20ms的时间里可以并发十个请求，虽然是每秒允许五个，但这种短时间内最大并发十个可能
并不是太好的选择。
于是我问了Ai，Ai提到了令牌桶算法，虽然名字很抽象，但是概念并不是很难理解，简单来说就是
根据pass time 与最大阈值来阶段性的生成可允许访问的数量，用代码抽象出来的话就是这样的

1. export class Bucket {
2.   private capacity: number; // 桶容量
3.   public current: number; // 当前水量
4.   private interval: number; // 补水速率
5.   private lastTime: number; // 补水时间
7.   constructor(capacity: number, interval: number) {
8.     this.capacity = capacity;
9.     this.current = capacity;
11.     this.interval = interval;
12.     this.lastTime = Date.now();
13.   }
14. ｝

复制代码

这里最关键的是理解capacity和interval字段的含义，capacity就是充当了阈值，而interval代表了“流速”（生成令牌的速度），比如一秒允许五次，那么0.8秒就生成了 5*0.8 = 4 个令牌，而current代表当前的令牌量（水桶中水的体积）,
那么用户每次在接口的访问，就会消耗一个令牌，然后在消耗令牌之前会先进行补充，因为每次请求之间理论上是有时间间隔的，哪怕很小，所以，需要先补充这段时间间隔产生的令牌量，有可能是零点几个，取决于接口并发访问的时间
所以这里来把方法定义一下

1.   refill() {
2.     const now = Date.now();
3.     const time = Number(((now - this.lastTime) / 1000).toFixed(1));
4.     const water = time * this.interval; // 需要补水量
5.     console.log("water:", water);
6.     this.current = Math.min(this.current + water, this.capacity);
7.     console.log("current:", this.current);
9.     this.lastTime = now;
10.   }
12.   /**
13.    * @description 消耗
14.    */
15.   consume() {
16.     this.refill();
17.     if (this.current >= 1) {
18.       this.current--;
19.       return true;
20.     }
21.     return false;
22.   }

复制代码

 
这里consume方法返回值，是因为我要根据用户来区分桶，不同的用户有不同的桶，所以我还要再抽象一个用户控制的类，这里采用内存存储，在生产环境中，可以使用db数据库或者redis缓存数据库来存储，当然为了方便我这里采用内存来达到最小可行性的目的

1. import { NextFunction, Response, Request } from "express";
3. import { Bucket } from "./Bucket";
5. class UserCircumScribe {
6.   static UserData: Map<string, Bucket> = new Map();
8.   middleware(req: Request, res: Response, next: NextFunction) {
9.     const userId = req.header("user-id") as string;
10.     if (!userId)
11.       return res.status(400).json({
12.         code: 400,
13.         message: "user-id is required",
14.       });
16.     if (!UserCircumScribe.UserData.has(userId)) {
17.       UserCircumScribe.UserData.set(userId, new Bucket(1, 0.1));
18.     }
19.     const userBucket = UserCircumScribe.UserData.get(userId);
20.     if (userBucket) {
21.       const pass = userBucket.consume();
22.       if (pass) {
23.         return next();
24.       }
25.       res.status(429).json({
26.         code: 429,
27.         message: "请求次数超出限制",
28.       });
29.     }
30.   }
31. }
33. export const circumscribe = new UserCircumScribe();

复制代码

因为我使用的是express.js框架，那么我就可以利用express.js的中间件机制来加载到某个接口或者全部接口上
// app.ts

1. import express, { Request, Response } from "express";
2. import { circumscribe } from "./lib/UserCircumscribe";
4. const app = express();
6. app.get("/", (req: express.Request, res: express.Response) => {
7.   res.send("Hello World!");
8. });
10. app.get(
11.   "/restriction",
12.   circumscribe.middleware,
13.   (_: Request, res: Response) => {
14.     res.status(200).json({
15.       message: "一些数据",
16.     });
17.   }
18. );
20. app.listen(3000);

复制代码

OK，我这里为了方便测试，设置的桶容量默认是1，流速是0.1，意思是十秒钟才生成一个令牌，看看效果

 

这样的话，一个令牌桶算法就在node.js中实现完成了，不过express生态中有提供的现成的包可以用：
express-rate-limit - npm
 

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！