为Gin设计中间件(三)：限流

一、背景

痛点：

1. 对于 登录和注册模块 ，最明显的漏洞是：”任何人都可以注册，任何人都可以登录“，也就是说，万一有一个人用 shell 脚本拼命给你发注册请求、登录请求，系统负载就会很高。
2. 若用限流，怎么标识某个用户？设定怎样的阈值呢？被限流的请求怎么办？
3. Gin 有许多开源的限流插件，但存在一定程度上的并发问题。

介绍：

1. 限流，限制每一个用户，每秒最多发送固定数量的请求。
2. web端用 IP 表示一个用户，即限流针对的是 IP（APP端用设备序列号会更好）。当然，在使用 IP 的情况下，我们可能会误把不同的人看成 是同一个人。但是只要我们限制的阈值不是很小，就不会有问题。
3. 理论上来说，阈值应该是通过压测来得到的。比如说你压测整个系统，发现最多只能撑住每秒 1000 个请求，那么阈值就是 1000 。而我们是针对个人，搞不了压测。所以可以凭借经验来设置，比如说我们正常人手速，一秒钟撑死一个请求，那么 就算我们考虑到共享 IP 之类的问题，给个每秒 100 也已经足够了。
4. 只能拒绝被限流的请求，也就是返回错误。 这个错误，不同公司有不同的规范。如果你自己决策的话，可以返回什么服务器繁忙之类的信息。

二、实现

（1）builder.go

package ratelimit

import (
    _ "embed"
    "fmt"
    "github.com/gin-gonic/gin"
    "github.com/redis/go-redis/v9"
    "log"
    "net/http"
    "time"
)

type Builder struct {
    prefix   string
    cmd      redis.Cmdable
    interval time.Duration
    // 阈值
    rate int
}

//go:embed slide_window.lua
var luaScript string

func NewBuilder(cmd redis.Cmdable, interval time.Duration, rate int) *Builder {
    return &Builder{
       cmd:      cmd,
       prefix:   "ip-limiter",
       interval: interval,
       rate:     rate,
    }
}

func (b *Builder) Prefix(prefix string) *Builder {
    b.prefix = prefix
    return b
}

func (b *Builder) Build() gin.HandlerFunc {
    return func(ctx *gin.Context) {
       limited, err := b.limit(ctx)
       if err != nil {
          log.Println(err)
          // 这一步很有意思，就是如果这边出错了
          // 要怎么办？
          ctx.AbortWithStatus(http.StatusInternalServerError)
          return
       }
       if limited {
          log.Println(err)
          ctx.AbortWithStatus(http.StatusTooManyRequests)
          return
       }
       ctx.Next()
    }
}

func (b *Builder) limit(ctx *gin.Context) (bool, error) {
    key := fmt.Sprintf("%s:%s", b.prefix, ctx.ClientIP())
    return b.cmd.Eval(ctx, luaScript, []string{key},
       b.interval.Milliseconds(), b.rate, time.Now().UnixMilli()).Bool()
}

（2）slide_window.lua 【使用 lua 脚本是为了避免并发问题】

-- 1, 2, 3, 4, 5, 6, 7 这是你的元素
-- ZREMRANGEBYSCORE key1 0 6
-- 7 执行完之后

-- 限流对象
local key = KEYS[1]
-- 窗口大小
local window = tonumber(ARGV[1])
-- 阈值
local threshold = tonumber( ARGV[2])
local now = tonumber(ARGV[3])
-- 窗口的起始时间
local min = now - window

redis.call('ZREMRANGEBYSCORE', key, '-inf', min)
local cnt = redis.call('ZCOUNT', key, '-inf', '+inf')
-- local cnt = redis.call('ZCOUNT', key, min, '+inf')
if cnt >= threshold then
    -- 执行限流
    return "true"
else
    -- 把 score 和 member 都设置成 now
    redis.call('ZADD', key, now, now)
    redis.call('PEXPIRE', key, window)
    return "false"
end

（3）在 middleware 中调用

func InitGinMiddlewares(hdl ijwt.Handler, l logger.LoggerV1) []gin.HandlerFunc {
    return []gin.HandlerFunc{
    // note 限流ratelimit.NewBuilder(redis.NewClient(&redis.Options{
         Addr: "localhost:6379",
    }), time.Second, 100).Build(),
}

三、总结

基于 Redis 实现限流：考虑到整个单体应用部署多个实例，用户的请求经过负载均衡之类的东西之后，就不一定落到同 一个机器上了，因此需要 Redis 来计数。