在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