掘金 后端 ( ) • 2024-06-30 17:58

一、背景

痛点

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

image.png 介绍

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

image.png

二、实现

(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 中调用

// note 限流
func(ctx *gin.Context) {
    redisClient := redis.NewClient(&redis.Options{
       Addr: "localhost:6379",
    })
    ratelimit.NewBuilder(redisClient, time.Second, 100).Build()
},

三、总结

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

image.png