Redis-based bucket rate limiting

fetch-tokens.lua

local bucket_key = KEYS[1]
local requested  = tonumber(ARGV[1])

local time   = redis.call('TIME')
local now    = tonumber(time[1]) * 1000000 + tonumber(time[2]) -- time in microseconds
local bucket = redis.call('HMGET', bucket_key, 'tokens', 'next_refill')

local tokens, next_refill = bucket[1], bucket[2]

if (not tokens) then
    -- initialize with defaults if token bucket doesn't yet exist
    local capacity = 1000
    tokens         = capacity
    redis.call('HMSET', bucket_key,
                'capacity',      capacity,
                'refill_amount', 10,           -- refill 10...
                'last_refill',   now,
                'next_refill',   now + 100000) -- every 100 ms, so 100/sec
else
    tokens, next_refill = tonumber(tokens), tonumber(next_refill)
    if (now >= next_refill) then
        bucket = redis.call('HMGET', bucket_key, 'capacity', 'refill_amount', 'last_refill')
        local capacity, refill_amount, last_refill = tonumber(bucket[1]), tonumber(bucket[2]), tonumber(bucket[3])

        local refill_interval = next_refill - last_refill
        local num_periods     = math.floor((now - last_refill) / refill_interval)
        tokens                = math.min(capacity, tokens + num_periods * refill_amount)
        last_refill           = last_refill + num_periods * refill_interval
        next_refill           = last_refill + refill_interval

        -- intentional code duplication and early return here to call redis only once on this code path.
        local granted = math.min(requested, tokens)
        tokens        = tokens - granted
        redis.call('HMSET', bucket_key, 'tokens', tokens, 'last_refill', last_refill, 'next_refill', next_refill)
        return granted
    end
end

local granted = math.min(requested, tokens)
tokens        = tokens - granted
redis.call('HSET', bucket_key, 'tokens', tokens)
return granted

init-bucket.lua

local bucket = KEYS[1]

local total_capacity  = tonumber(ARGV[1])
local refill_amount   = tonumber(ARGV[2])
local refill_interval = tonumber(ARGV[3]) -- in microseconds (i.e. redis precision)
local time = redis.call('TIME')
local now = tonumber(time[1]) * 1000000 + tonumber(time[2]) -- time in microseconds

redis.call('HMSET', bucket,
            'tokens',        total_capacity,
            'capacity',      total_capacity,
            'refill_amount', refill_amount,
            'last_refill',   now,
            'next_refill',   now + refill_interval)

return redis.status_reply('OK')