基于Redis实现一个简单的固定窗口限流器

大家好，我是渔夫子。

限流器是在大流量中保护服务资源的一种常用手段。限流器的实现有令牌桶方式、固定窗口限流器和滑动窗口限流器。本文介绍了基于Redis如何快速的实现固定窗口限流器。

最近在我们的项目中需要快速的实现一个流量限流器，而目前项目中已经有在用Redis了。

固定窗口限流器：它是在固定的时间窗口（例如一分钟）内计算接收到的请求数量。一旦达到最大请求数量，额外的请求将被拒绝，直到下一个窗口开始。

要基于Redis实现固定窗口限流器非常简单，如下lua代码：

local current
current = redis.call("INCR", KEYS[1])
if tonumber(current) == 1 then
 redis.call("EXPIRE", KEYS[1], 60)
end
return current

每次运行这个脚本时，它都会获取一个键并将其值递增1。如果是第一次递增该键时，都会设置一个60秒的过期时间。它返回递增后的当前值。

该键在首次设置60秒后过期。一旦过期，它将在下一个请求时再次设置。

当服务收到一个请求时，就可以调用该段代码。如果脚本返回的值大于允许的值，则由于速率限制而中止该请求。如果返回的值不大于允许的值，则处理该请求。

const script = `
local current
current = redis.call("INCR", KEYS[1])
if tonumber(current) == 1 then
 redis.call("EXPIRE", KEYS[1], 60)
end
return current
`

func isRateLimited(ctx context.Context, key string, limit int64) (bool, error) { 
 v, err := redisClient.Eval(ctx, script, []string{key}).Result()
 if err != nil {
  return false, err
 }
 n, _ := v.(int64)
 return n > int64(limit), nil
}

isRateLimited函数可以按如下方式使用：

func handleLogin(r *http.Request, w http.ResponseWriter) {
 username := r.FormValue("username")

 limited, _ := isRateLimited(context.TODO(), fmt.Sprintf("rateLimit:login:username:%s", username, 5))
 if limited {
  http.Error(w, "Too Many Attempts", http.StatusTooManyRequests)
  return
 }

 // ...
}

这样就可以工作了。

请注意，固定窗口限流器虽然可以有效抵御持续攻击，但可能会影响合法用户的体验。

在上面的示例中，我们基于在登录流程中使用的用户名进行速率限制。如果是基于其他指标进行限流（例如传入请求的远程IP地址），那么该限流器是不起作用的。

特别说明：你的关注，是我写下去的最大动力。点击下方公众号卡片，直接关注。关注送《100个go常见的错误》pdf文档、经典go学习资料。