Token bucket rate limiting: algorithm and Redis implementation

An API endpoint is being hammered by a client, consuming all server resources and degrading service for other users.

Implement token bucket rate limiting using Redis with atomic Lua execution:

import { createClient } from 'redis';
const redis = createClient();

async function tokenBucketAllow(key, { capacity = 10, refillRate = 1, refillIntervalMs = 1000 } = {}) {
  const now = Date.now();
  const bucketKey = `ratelimit:${key}`;

  // Lua script executes atomically on the Redis server
  const script = `
    local tokens = tonumber(redis.call('hget', KEYS[1], 'tokens'))
    local last = tonumber(redis.call('hget', KEYS[1], 'last'))
    local now = tonumber(ARGV[1])
    local capacity = tonumber(ARGV[2])
    local refill_rate = tonumber(ARGV[3])
    local interval = tonumber(ARGV[4])

    if not tokens then
      tokens = capacity
      last = now
    end

    local elapsed = math.floor((now - last) / interval)
    tokens = math.min(capacity, tokens + elapsed * refill_rate)
    last = last + elapsed * interval

    if tokens < 1 then
      redis.call('hset', KEYS[1], 'tokens', tokens, 'last', last)
      redis.call('expire', KEYS[1], 3600)
      return 0
    end

    redis.call('hset', KEYS[1], 'tokens', tokens - 1, 'last', last)
    redis.call('expire', KEYS[1], 3600)
    return 1
  `;

  const allowed = await redis.sendCommand(['FCALL', script, '1', bucketKey, String(now), String(capacity), String(refillRate), String(refillIntervalMs)]);
  return allowed === 1;
}

// Express middleware
app.use('/api/', async (req, res, next) => {
  const allowed = await tokenBucketAllow(req.ip, { capacity: 60, refillRate: 1, refillIntervalMs: 1000 });
  if (!allowed) {
    res.setHeader('Retry-After', '1');
    return res.status(429).json({ error: 'Rate limit exceeded' });
  }
  next();
});

Using an atomic Lua script in Redis makes the check-and-decrement operation race-condition free. Token bucket allows bursting up to the capacity before throttling, which accommodates legitimate usage spikes.