Golang 利用 redis 实现每秒请求限量

限量核心代码：limit.go

type Limit struct {
    Name    string
    Key     string
    Rate    int64
    Max     int64
    Default int64
}

func (l *Limit) Add(a, b float64) float64 {
    return a + b
}

// 令牌 Unit is `qps/seconds`
// rate 允许多大的请求访问，请求数/秒
// max 最大的请求数
// defaults 初始化请求数
// return 1验证通过，0被限量
func (l *Limit) getLimit(config *Limit) int {
    name := config.Name
    key := config.Key

    now := time.Now().Unix()
    sKey := GetStorekey(name, key)
    nKey := GetNextTimeKey(name, key)

    rate := config.Rate
    max := config.Max
    defaults := config.Default

    if luaScript == "" {
        logs.Error("load lua script error")
        return 0
    }

    result2, err := common.RedisLimitClient.ScriptLoad(luaScript).Result() //返回的脚本会产生一个sha1哈希值,下次用的时候可以直接使用这个值
    if err != nil {
        logs.Error("fillLimit-GetLimit-ScriptLoad-err", err)
        return 0
    }
    foo := common.RedisLimitClient.EvalSha(result2, []string{sKey, nKey}, now, rate, max, defaults)
    result, err := foo.Int()
    if err != nil {
        logs.Error("fillLimit-GetLimit-EvalSha-err", err)
        return 0
    }
    return result

}

使用的lua脚本: limit.lua

local sKey = KEYS[1];
local nKey = KEYS[2];
local now = tonumber(ARGV[1]);
local rate = tonumber(ARGV[2]);
local max = tonumber(ARGV[3]);
local default = tonumber(ARGV[4]);

local sNum = redis.call('get', KEYS[1]);
if((not sNum) or sNum == nil)
then
    sNum = 0
end

sNum = tonumber(sNum);

local nNum = redis.call('get', KEYS[2]);
if((not nNum) or nNum == nil)
then
    nNum = now
    sNum = default
end

nNum = tonumber(nNum);

local newPermits = 0;
if(now > nNum)
then
    newPermits = (now-nNum)*rate+sNum;
    sNum = math.min(newPermits, max)
end

local isPermited = 0;
if(sNum > 0)
then
    sNum = sNum -1;
    isPermited = 1;
else
    sNum = 0;
    isPermited = 0;
end

redis.call('set', KEYS[1], sNum);
redis.call('set', KEYS[2], now);

return isPermited

初始化时加载 lua脚本：loadLua.go

var luaScript = ""

func init() {
    loadLua()
}

func loadLua() {
    luaPath := "models/fill/limit/limit.lua"
    luaFile, err := filepath.Abs(luaPath)
    if err != nil {
        logs.Error("fill-limit.lua", err)
        return
    }
    luaScriptByte, err := os.ReadFile(luaFile)
    if err != nil {
        logs.Error("fill-ReadFile-limit.lua", err)
        return
    }
    luaScript = string(luaScriptByte)
}

定义redis中使用的key：keys.go

// GetLimitKey 默认限流key
func GetLimitKey() string {
    return "{fill:limit}"
}

// GetStorekey 获取令牌桶key
func GetStorekey(name string, key string) string {
    return fmt.Sprintf("%s:%s:store", name, key)
}

func GetNextTimeKey(name string, key string) string {
    return fmt.Sprintf("%s:%s:next", name, key)
}

入口测试：main.go

func main() {
    // 限流
    config := &Limit{
        Name:    GetLimitKey(),
        Key:     "suibian",
        Rate:    1,
        Max:     6,
        Default: 5,
    }
    for i := 0; i < 10; i++ {
        result := config.GetLimit(config)
        println(result)
    }
}