手写一个基于nestjs的限流器

前言

限流是大型系统必备的保护措施，常用的限流算法主要有固定时间窗口，滑动时间窗口，漏桶，令牌桶等。本文将会写道的方案是使用滑动时间窗口算法，通过拒绝请求的方式来达到限流的目的。本文的实现方式是 redis , lua 脚本 以及 Nestjs Guard 来实现限流的效果。

概念浅析

这里简单说一下固定时间窗口和滑动时间窗口的概念

固定时间窗口 它可以解决每时间单位（可以是秒或者分钟等等），允许访问的次数。但是无法控制频率。举例1分钟允许访问100 次，可能前10 秒访问了90次，后面只有10次机会了。还有一个问题就是在两个时间单位的临界值上可能会超出阈值，继续用前面的例子，第59秒访问了60次，第二个时间单位前10秒访问了50 次，在横跨两个时间单位的20秒中，超出了阈值 (110>100)

滑动时间窗口 可以改善固定窗口的所带来超出阈值的问题。它将每个单位之间分割成若干小周期，当前时间单位不再是固定的，而是根据当前请求时间往后移动，即所谓滑动窗口。每个周期分的越小，限流控制的越精细。

具体实现

使用的主要包的版本 nestjs 8.0.0 ioredis 5.3.2

我们主要实现以下几个东西

一个 guard 文件用于实现限流的业务逻辑
一个 decorator文件 , 装饰器，用于设置当前接口限流的频率，允许访问次数等字段
一个 redis 类和一个lua 脚本

redis 相关

主要就是通过lua 脚本进行计数，达到限流的目的。这里做了一个优化，对执行lua 取了hash 值，在redis 运行一次后，可以使用evalsha 直接运行脚本，避免二次载入脚本。

ts 复制代码

import { Injectable } from '@nestjs/common';
import * as Redis from 'ioredis';
import { ConfigService } from '@app/common';
import { createHash } from 'crypto'
import { v4 as uuidv4 } from 'uuid';

const rateLimitScript = ""// 后面单独列出

@Injectable()
export class RedisService {
    private readonly redisClient: Redis.Redis;
    private luaScript: any;
    constructor(
        private readonly configService: ConfigService,
    ) {

        const self = this;
        const connConfig = this.configService.get("redisService")
        this.redisClient = new Redis.Redis(connConfig)
        this.luaScript = {
            rateLimit: {
                script: rateLimitScript,
                hash: self.hashStr(rateLimitScript)
            },
        }
        
    }

    private hashStr(value: string) {
        return createHash("sha1").update(value).digest('hex')
    }

    async rateLimit(opts: any): Promise<boolean> {

        const { key, limit, windowSize } = opts;
        const uuid = uuidv4()
        let result;
        const { script, hash } = this.luaScript.rateLimit
  
        try {

            const shaResult = await this.redisClient.evalsha(hash, 1, key, limit, windowSize, uuid)
            result = shaResult

        } catch (error) {

            const shaResult = await this.redisClient.eval(script, 1, key, limit, windowSize, uuid)
            result = shaResult

        }
        return result == 1
    }
}

接下来展示lua 脚本

lua 复制代码

--传入四个参数 分别是key,限制次数,时间范围,唯一值
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local windowSize = tonumber(ARGV[2]) --单位毫秒
local uuid = ARGV[3] -- 唯一值是为了防止zset 重复

-- 使用redis 来获取时间，防止多进程生成相似的边界导致超频。时间单位是微秒
local date = redis.call("time")
local now = tonumber(date[1]) * 1000000 + tonumber(date[2])
local startTime = now - windowSize * 1000
local endTime = now +  1000000

-- 计算过期时间 时间单位是秒
local expireSec = tonumber(math.ceil(windowSize / 1000)) + 1

-- 统计当前zset数组里的数据,超出范围则返回0,
-- 否则做3件事,然后返回1
-- 1、向数组里增加新值
-- 2、删除数组中开始时间之前的数据,防止数组过大
-- 3、给数组续过期时间
local count = tonumber(redis.call('zcount', key, startTime, endTime))

if count + 1 > limit then
    return 0
else
    redis.call('zadd', key, now, uuid)
    redis.call('zremrangebyscore', key, 0, startTime - 100000)
    redis.call('expire', key, expireSec)
    return 1
end

装饰器相关

这个很简单就是，设置一下redis 键值的前缀，允许访问的次数和单位之间的长度。在这里设置了之后可以在 guard 里通过反射拿到这些值

ts 复制代码

import { SetMetadata } from '@nestjs/common';

export interface rateLimitOptions {
    keyPrefix: string,
    limit: number,
    windowSize: number
}

export const RateLimit = (options: rateLimitOptions): MethodDecorator => SetMetadata('rateLimit', options)

guard 相关

guard 就是把之前的部分整合了一下，如果当前接口没有设置限流参数则启用默认参数，keyprefix 取当前接口的路径。

ts 复制代码

import { Injectable, ExecutionContext, CanActivate } from "@nestjs/common";
import { Reflector } from '@nestjs/core';
import { RedisService, rateLimitOptions } from "@app/common";
import { BusinessException } from "@app/common";

@Injectable()
export class RateLimitGuard implements CanActivate {
    constructor(
        private reflector: Reflector,
        private redisService: RedisService
    ) { }

    private getIpFromRequest(request: { ip: string }): string {
        return request.ip?.match(/\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b/)?.[0]
    }

    async canActivate(context: ExecutionContext) {

         // 通过反射拿到前面设置的值
        const rateLimitConfig = this.reflector.get<rateLimitOptions>("rateLimit", context.getHandler());
           if (!rateLimitConfig) {
			 // 当前接口如果没设置参数则定义默认参数
            const cMethod = this.reflector.get("method", context.getHandler());// 是GET,POST 等http method
            const cPath = this.reflector.get("path", context.getHandler());// 接口的具体路径
            rateLimitConfig = {
                keyPrefix: cMethod + ":" + cPath,
                limit: 1,
                windowSize: 5000
            }
        }

        const { keyPrefix, limit, windowSize } = rateLimitConfig
        const request = context.switchToHttp().getRequest();
        const ip = this.getIpFromRequest(request)
        const key = keyPrefix + ":" + ip

        const isPass = await this.redisService.rateLimit({
            key,
            limit,
            windowSize
        })
        if (!isPass) {
            // 返回自定义的错误
            throw new BusinessException("RATE_LIMIT_EXCEEDED_LIMIT")
        }
        return true
    }
}

使用方法

引入guand 和RateLimit 装饰器，可以给特定路由增加限流保护

ts 复制代码

@Controller('user')
export class UserController {
  constructor(private readonly userService: UserService) { }
  
  @Public()
  @RateLimit({ keyPrefix: "login", limit: 3, windowSize: 1000 })
  @UseGuards(RateLimitGuard)
  @Post('register')
  register(@Body() createUserDto: CreateUserDto) {
    return this.userService.register(createUserDto);
   }
}

或者基于模块的也可以，这样路由里的就可以省略了，如果某些接口没设置RateLimit 参数，guard 内部就会使用默认统一参数。

ts 复制代码

@Module({
  providers: [
    {
      provide: APP_GUARD,
      useClass: RateLimitGuard
    }
  ],
})

使用ab 测试一下结果，为了便于测试设置为每5秒可以请求3次。用 ab进行两次测试，结果如下

shell 复制代码

2023-11-25 17:50:39 - error - HttpExceptionFilter - d1385d48-183c-4fbf-b751-4d0b6786f5ba : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:39 - error - HttpExceptionFilter - 65f0e427-92e4-4854-a6cc-116c70daac61 : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:39 - error - HttpExceptionFilter - b24b8f7e-f961-45d1-a909-36219fc5d112 : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:39 - error - HttpExceptionFilter - 9c65d452-8eeb-4c40-a76e-b5bf01524ebb : {"validatorCode":30000,"validatorMessage":"请求频率过快"} - {}
2023-11-25 17:50:39 - error - HttpExceptionFilter - 1065a900-bb55-4514-9b55-08cc57509e37 : {"validatorCode":30000,"validatorMessage":"请求频率过快"} - {}
2023-11-25 17:50:45 - error - HttpExceptionFilter - a8176f1c-2788-4e2a-8267-4e2ffadb6238 : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:45 - error - HttpExceptionFilter - db911d44-ee8b-4da2-a87b-fa0bcb433c45 : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:45 - error - HttpExceptionFilter - 3c59e335-441b-4907-80e4-0d807e5bfb01 : {"validatorCode":10005,"validatorMessage":"用户已存在"} - {}
2023-11-25 17:50:45 - error - HttpExceptionFilter - 6f48108b-bc8e-4fb6-8231-fa5a9cd22b5f : {"validatorCode":30000,"validatorMessage":"请求频率过快"} - {}
2023-11-25 17:50:45 - error - HttpExceptionFilter - 7c77c1df-bfe6-4f72-b75a-0d9a1211fe64 : {"validatorCode":30000,"validatorMessage":"请求频率过快"} - {}

达到要求，收工。

手写一个基于nestjs的 限流器

前言

概念浅析

具体实现

redis 相关

装饰器相关

guard 相关

使用方法

手写一个基于nestjs的限流器