redis-token-bucket-ratelimiter/lua/rollingLimit.lua.json

Rolling rate limit in redis via a lua script

{"script":"-- valueKey timestampKey | limit intervalMS nowMS [amount]\nlocal valueKey     = KEYS[1] -- \"limit:1:V\"\nlocal timestampKey = KEYS[2] -- \"limit:1:T\"\nlocal limit      = tonumber(ARGV[1])\nlocal intervalMS = tonumber(ARGV[2])\nlocal amount     = math.max(tonumber(ARGV[3]), 0)\nlocal force      = ARGV[4] == \"true\"\n\nlocal lastUpdateMS\nlocal prevTokens\n\n-- Use effects replication, not script replication;; this allows us to call 'TIME' which is non-deterministic\nredis.replicate_commands()\n\nlocal time = redis.call('TIME')\nlocal nowMS = math.floor((time[1] * 1000) + (time[2] / 1000))\nlocal initialTokens = redis.call('GET',valueKey)\nlocal initialUpdateMS = false\n\n\nif initialTokens == false then\n   -- If we found no record, we temporarily rewind the clock to refill\n   -- via addTokens below\n   prevTokens = 0\n   lastUpdateMS = nowMS - intervalMS\nelse\n   prevTokens = initialTokens\n   initialUpdateMS = redis.call('GET',timestampKey)\n\n   if(initialUpdateMS == false) then -- this is a corruption\n      -- we make up a time that would fill this limit via addTokens below\n      lastUpdateMS = nowMS - ((prevTokens / limit) * intervalMS)\n   else\n      lastUpdateMS = initialUpdateMS\n   end\nend\n\n-- tokens that should have been added by now\n-- note math.max in case this ends up negative (clock skew?)\n-- now that we call 'TIME' this is less likely to happen\nlocal addTokens = math.max(((nowMS - lastUpdateMS) / intervalMS) * limit, 0)\n\n-- calculated token balance coming into this transaction\nlocal grossTokens = math.min(prevTokens + addTokens, limit)\n\n-- token balance after trying this transaction\nlocal netTokens = grossTokens - amount\n\n-- time to fill enough to retry this amount\nlocal retryDelta = 0\n\nlocal rejected = false\nlocal forced = false\n\nif netTokens < 0 then -- we used more than we have\n   if force then\n      forced = true\n      netTokens = 0 -- drain the swamp\n   else\n      rejected = true\n      netTokens = grossTokens -- rejection doesn't eat tokens\n   end\n   -- == percentage of `intervalMS` required before you have `amount` tokens\n   retryDelta = math.ceil(((amount - netTokens) / limit) * intervalMS)\nelse -- polite transaction\n   -- nextNet == pretend we did this again...\n   local nextNet = netTokens - amount\n   if nextNet < 0 then -- ...we would need to wait to repeat\n      -- == percentage of `invervalMS` required before you would have `amount` tokens again\n      retryDelta = math.ceil((math.abs(nextNet) / limit) * intervalMS)\n   end\nend\n\n-- rejected requests don't cost anything, we'll wait for a costly request to update our values\n-- forced requests show up here as !rejected, but with netTokens = 0 (drained)\nif rejected == false then\n\n   redis.call('PSETEX',valueKey,intervalMS,netTokens)\n\n   if addTokens > 0 or initialUpdateMS == false then\n      -- we filled some tokens, so update our timestamp\n      redis.call('PSETEX',timestampKey,intervalMS,nowMS)\n   else\n      -- we didn't fill any tokens, so just renew the timestamp so it survives with the value\n      redis.call('PEXPIRE',timestampKey,intervalMS)\n   end\nend\n\nreturn { netTokens, rejected, retryDelta, forced }\n","sha1":"ed028bce56b447a97d9277e51b27346c681443e7"}