lua实现令牌桶限流算法

-- EVAL "$(cat tokenbucket.lua)" 1 keys1   arg1    arg2  arg3 arg4  arg5
-- EVAL "$(cat tokenbucket.lua)" 1 monName curTime level rate burst prisonTime

local PRISONSEC  = tonumber(ARGV[5]);
local SEC2USEC   = 1000000;
local PRISONUSEC = SEC2USEC * PRISONSEC;

-- 计算[lastTime, curTime] 这段时间内需要产生多少token.
local function generateTokens(lastTime, curTime, rate)
    local deltaTime = curTime - lastTime;
    if (deltaTime <= 0) then
        return 0;
    end
    return (deltaTime / SEC2USEC * rate);
end

-- 关监狱结束, 1: 是; 0: 否;
local function finishPrison(enterTime, curTime)
    local deltaTime = (curTime - enterTime);
    if (deltaTime < 0) or (deltaTime < PRISONUSEC) then
        return 0;
    end
    return 1;
end

-- token bucket 算法
local function tokenBucket()
    local curTime = tonumber(ARGV[1]);
    local level   = tonumber(ARGV[2]);
    local rate    = tonumber(ARGV[3]);
    local bucketCapacity = math.max(tonumber(ARGV[4]), 1.0);
    local curToken  = 0;
    local lastTime  = 0;
    local enterTime = 0;
    local content   = redis.call("HGET", KEYS[1], level);

    if (content) then
        curToken,lastTime,enterTime = string.match(content, "([0-9|.]+) (%d+) (%d+)");
        curToken = tonumber(curToken);
        lastTime = tonumber(lastTime);
        enterTime = tonumber(enterTime);

        -- in prison or not
        if (enterTime > 0) then
            if (finishPrison(enterTime, curTime) == 1) then
                curToken = curToken + generateTokens(lastTime, curTime, rate);
                lastTime = curTime;
                enterTime = 0;
            else
                enterTime = curTime;
            end
        else
            curToken = curToken + generateTokens(lastTime, curTime, rate);
        end
    else
        curToken  = bucketCapacity;
        lastTime  = curTime;
        enterTime = 0;
    end

    if (enterTime == 0) then
        lastTime = curTime;
        if (curToken < 1) then
            enterTime = curTime;
        else
            curToken = math.min(curToken, bucketCapacity) - 1;
        end
    end
    local item = string.format("%f %.0f %.0f", curToken, lastTime, enterTime);
    redis.call("HSET", KEYS[1], level, item);
    redis.call("EXPIRE", KEYS[1], 2 * 24 * 60 * 60);
    return item;
end

-- main:
return tokenBucket();