Rate Limit Algorithms

1// Executed server-side in the store, so read-refill-decrement-write
2// cannot interleave with another instance doing the same thing.
3// TIME comes from the store: instances with skewed clocks would
4// otherwise credit different amounts of refill.
5const TOKEN_BUCKET = `
6  local bucket   = redis.call('HMGET', KEYS[1], 'tokens', 'updated')
7  local capacity = tonumber(ARGV[1])   -- burst we are willing to absorb
8  local refill   = tonumber(ARGV[2])   -- tokens per second = sustained rate
9  local cost     = tonumber(ARGV[3])   -- cost-weighted: not every call is 1
10  local now      = tonumber(redis.call('TIME')[1])
11 
12  local tokens  = tonumber(bucket[1]) or capacity
13  local updated = tonumber(bucket[2]) or now
14 
15  -- refill for elapsed time, never above capacity
16  tokens = math.min(capacity, tokens + (now - updated) * refill)
17 
18  if tokens < cost then
19    -- how long until enough tokens exist: a real Retry-After,
20    -- not a constant we guessed
21    local wait = math.ceil((cost - tokens) / refill)
22    return { 0, tokens, wait }
23  end
24 
25  tokens = tokens - cost
26  redis.call('HMSET', KEYS[1], 'tokens', tokens, 'updated', now)
27  redis.call('EXPIRE', KEYS[1], math.ceil(capacity / refill) * 2)
28  return { 1, tokens, 0 }
29`
30 
31async function allow(key: string, cost = 1) {
32  const [ok, remaining, retryAfter] = await redis.eval(
33    TOKEN_BUCKET, 1, `rl:${key}`, CAPACITY, REFILL_PER_SEC, cost,
34  )
35  return { allowed: ok === 1, remaining, retryAfter }
36}

Three things are doing work beyond the arithmetic. The store's clock removes instance skew from the refill calculation. The expiry bounds key growth so the limiter cannot leak memory across every key it has ever seen. And cost is what lets one expensive endpoint charge more than one cheap one against the same bucket.