1.令牌桶限流的思想

以恒定的速率去生成令牌,存放到一定容量的令牌桶中。桶中存放的令牌数量不会超过桶的容量。从客户端进来的请求需要去获取令牌,只有拿到令牌的请求才会被处理,否则被限流。

2.lua脚本

-- 限流函数
local function limit_rate(key, limit, interval)
    -- 获取当前时间戳
    local now = redis.call('TIME')[1]
 
    -- 获取令牌桶的信息
    local rate_limit_info = redis.call('HMGET', key, 'last_time', 'tokens')
    local last_time = tonumber(rate_limit_info[1]) or 0
    local tokens = tonumber(rate_limit_info[2]) or limit
 
    -- 计算时间间隔内新增的令牌数
    local elapsed = now - last_time
    local refill = math.floor(elapsed * limit / interval)
 
    -- 重新填充令牌桶
    if refill > 0 then
        tokens = math.min(tokens + refill, limit)
        redis.call('HSET', key, 'tokens', tokens)
        redis.call('HSET', key, 'last_time', now)
    end
 
    -- 检查令牌是否足够
    if tokens > 0 then
        tokens = tokens - 1
        redis.call('HSET', key, 'tokens', tokens)
        return true
    else
        return false
    end
end
 
-- 调用限流函数
return limit_rate(KEYS[1], ARGV[1], ARGV[2])

解释:首先根据key从redis中获取相关的令牌桶配置(先获取一次,如果没有值的话就初始化),然后根据传入的当前时间去计算上次获取令牌到这次中间产生了多少个令牌(如果产生的令牌数加上剩余的超过令牌桶的容量的话,选择让剩余的令牌数量等于容量。)。然后更新redis中的相关的值。然后如果剩余的令牌数大于零的话,就通过此次的请求,让令牌数减一并更新redis,否则不通过。

3.go代码

package main
 
import (
	"fmt"
	"github.com/go-redis/redis/v8"
)
 
func main() {
	rdb := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // 如果有密码,需要填写
		DB:       0,  // Redis 数据库索引
	})
 
	// 调用限流 Lua 脚本
	script := `
		-- Lua 脚本内容
		...  -- 省略 Lua 脚本的内容
	`
	key := "my_bucket_key"
	limit := 10
	interval := 1
 
	result, err := rdb.Eval(script, []string{key}, limit, interval).Result()
	if err != nil {
		fmt.Println("执行限流 Lua 脚本出错:", err)
		return
	}
 
	if result == int64(1) {
		fmt.Println("限流通过")
	} else {
		fmt.Println("限流不通过")
	}
}

解释:通过go-redis的Client.Eval()函数来调用命令行执行脚本,并将参数传递过去(令牌桶配置在redis中的key、令牌桶的容量、令牌桶生成满令牌数量的时间间隔),拿到返回值,通过返回值判断是否被限流了并做下一步的响应。