github.com/rudderlabs/rudder-go-kit@v0.30.0/throttling/lua/sortedset.lua (about)

     1  --[[
     2  To debug this script you can add these entries here in the script and then check the Redis server output:
     3  redis.log(redis.LOG_NOTICE, "some label", some_variable)
     4  
     5  For more information please refer to: https://redis.io/docs/manual/programmability/lua-debugging/
     6  --]]
     7  
     8  local key = tostring(KEYS[1])
     9  local cost = tonumber(ARGV[1])
    10  local rate = tonumber(ARGV[2])
    11  local period = tonumber(ARGV[3]) * 1000 * 1000 -- converting to microseconds
    12  
    13  -- Check number of requests first
    14  if cost < 1 then
    15      -- nothing to do, the user didn't ask for any tokens
    16      return { 0, "" }
    17  end
    18  
    19  local current_time = redis.call('TIME')
    20  local microseconds = current_time[2]
    21  while string.len(microseconds) < 6 do
    22      -- in case the microseconds part (i.e. current_time[2]) is less than 6 digits
    23      microseconds = "0" .. microseconds
    24  end
    25  
    26  local current_time_micro = tonumber(current_time[1] .. microseconds)
    27  local trim_time = current_time_micro - period
    28  
    29  -- Remove all the requests that are older than the window
    30  redis.call('ZREMRANGEBYSCORE', key, 0, trim_time)
    31  
    32  -- Get the number of requests in the current window
    33  local used_tokens = redis.call('ZCARD', key)
    34  
    35  -- If the number of requests is greater than the max requests we hit the limit
    36  if (used_tokens + cost) > tonumber(rate) then
    37      local next_to_expire = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')[2]
    38      local retry_after = next_to_expire + period - current_time_micro
    39  
    40      return { current_time_micro, "", retry_after }
    41  end
    42  
    43  -- seed needed to generate random members in case of collision
    44  math.randomseed(current_time_micro)
    45  
    46  local members = ""
    47  for i = 1, cost, 1 do
    48      local added = 0
    49      local member = current_time_micro .. i
    50  
    51      while added == 0 do
    52          -- ZADD key score member
    53          -- current_time_micro is used for the "score" to sort the members in the sorted set
    54          -- "member" is the token representing the request and has to be unique (score does not have to be unique)
    55          added = redis.call('ZADD', key, current_time_micro, member)
    56          if added == 0 then
    57              -- the member already exists, we need to generate a new one
    58              -- the overhead for math.random is used only in case of collision
    59              member = current_time_micro .. i .. math.random(1, 1000000) -- using a string, no risk of truncation
    60          else
    61              members = members .. member .. ","
    62          end
    63      end
    64  end
    65  
    66  redis.call('EXPIRE', key, period)
    67  
    68  members = members:sub(1, -2) -- remove the last comma
    69  return {
    70      current_time_micro,
    71      members,
    72      0 -- no retry_after
    73  }