Redis入门12-Lua脚本

写了一个lua脚本用于，记录一下

1. IP请求每分钟500次
2. 每个用户最近10min只能发10次
3. 每个用户最近1小时只能发20次
4. 每个用户最近1天只能发30次

-- 设置限制
-- 参数说明：
-- key：限制的关键字，可以是用户ID或者IP地址
-- period：时间段，单位为秒，可以是 600(10分钟)、3600(1小时)、86400(1天)
-- limit：时间段内允许的最大请求次数
-- ip_limit：时间段内IP地址允许的最大请求次数
-- expire：超时时间，单位为秒
-- 参数校验函数，校验传入的参数是否合法
-- 参数校验方法
local function validateParams(period, limit, ip_limit, expire)
    if not (tonumber(period) == 600 or tonumber(period) == 3600 or tonumber(period) == 86400) then
        return false
    end
    if not (tonumber(limit) and tonumber(ip_limit) and tonumber(expire)) then
        return false
    end
    return true
end

local function setLimit(key, period, limit, ip_limit, expire)
  	-- 校验参数
  	if not validateParams(period, limit, ip_limit, expire) then
        return 1
    end
  
    local now = redis.call("TIME")[1]
    local user_key = "user:" .. key .. ":" .. period
    local ip_key = "ip:" .. key .. ":" .. period
		
    -- 获取用户和IP的访问次数以及超时时间
    local user_count = tonumber(redis.call("GET", user_key) or "0")
    local ip_count = tonumber(redis.call("GET", ip_key) or "0")
    local user_expire = tonumber(redis.call("PTTL", user_key) or "-1")
    local ip_expire = tonumber(redis.call("PTTL", ip_key) or "-1")

    -- 如果用户或IP的超时时间小于0，则将其重置为 expire
    if user_expire < 0 then
        redis.call("SETEX", user_key, expire, "0")
    end
    if ip_expire < 0 then
        redis.call("SETEX", ip_key, expire, "0")
    end

    -- 如果用户或IP的访问次数超过了限制，返回 0 表示限制
    if user_count >= limit or ip_count >= ip_limit then
        return 0
    end

    -- 计算时间段的开始和结束时间
    local period_start = now - (now % period)
    local period_end = period_start + period

    -- 使用事务操作，将用户和IP的访问次数增加 1，并设置超时时间
    redis.call("MULTI")
    redis.call("INCR", user_key)
    redis.call("EXPIREAT", user_key, period_end)
    redis.call("INCR", ip_key)
    redis.call("EXPIREAT", ip_key, period_end)

    -- 获取用户和IP的访问次数，如果任意一个访问次数超过了限制，则返回 0 表示限制
    local result = redis.call("EXEC")
    if tonumber(result[1]) > limit or tonumber(result[3]) > ip_limit then
        return 0
    end

    -- 计算用户和IP的访问次数总和，如果总和超过了每分钟限制，则设置用户和IP的超时时间为60秒
    local total_count = user_count + ip_count + 1
    if total_count > 500 then
        redis.call("EXPIRE", user_key, 60)
        redis.call("EXPIRE", ip_key, 60)
    end

    -- 返回 1 表示不限制
    return 1
end
-- 增加同步模式
redis.replicate_commands()
-- 调用 setLimit 函数
return setLimit(KEYS[1], ARGV[1], ARGV[2], ARGV[3], ARGV[4])

注意这里的 redis.replicate_commands() 目的从同步命令变更为同步内容，防止出现命令在主从设备不一致的情况(获取当前时间，当命令从主同步到从后由于时间差导致不一致)