Design a Distributed Rate Limiter

// Express middleware. 429 Too Many Requests is defined in RFC 6585;
// Retry-After is RFC 9110 section 10.2.3 and is the only header here that is
// actually standardised. The RateLimit-* fields are an IETF draft: widely
// copied, not a standard, so send them but do not rely on clients reading them.
import { consume } from './limiter.js';

export function rateLimit({ limit, windowMs, identify }) {
  return async function (req, res, next) {
    // `identify?.(req) ?? req.ip` is shorthand for the two longhand lines below.
    // Note that `??` falls back only on null/undefined, where `||` would also
    // fall back on an empty string - which a tenant id could legitimately be.
    const key = identify?.(req) ?? req.ip;
    // let key = identify == null ? undefined : identify(req);
    // if (key === null || key === undefined) key = req.ip;

    const { allowed, remaining, retryAfterMs } = await consume(key, { limit, windowMs });

    res.set('RateLimit-Limit', String(limit));
    res.set('RateLimit-Remaining', String(remaining));
    if (allowed) return next();

    // Retry-After is whole seconds. Round up: rounding down invites the client
    // straight back into another 429.
    res.set('Retry-After', String(Math.ceil(retryAfterMs / 1000)));
    return res.status(429).json({ error: 'rate_limited', retry_after_ms: retryAfterMs });
  };
}

// Two-tier limiter. The hot path is in-process; Redis is consulted once per
// block of permits rather than once per request. Use it when a single key is
// taking thousands of requests per second and the round trip is the bottleneck.
const BLOCK = 20;
const leases = new Map();   // key -> { permits, expiresAt }

export async function consumeBatched(key, opts) {
  const now = Date.now();
  const lease = leases.get(key);
  if (lease && lease.expiresAt > now && lease.permits > 0) {
    lease.permits -= 1;
    return { allowed: true, source: 'local' };
  }

  // Deduct the whole block from the global counter up front. Because the
  // permits are spent from the shared budget at lease time, the fleet cannot
  // exceed the limit. What it can do is distribute it badly: a process that
  // goes quiet sits on permits a busy process needed. That unfairness, not
  // over-admission, is the thing you are accepting here.
  const { allowed, granted } = await leaseBlock(key, { ...opts, want: BLOCK });
  // `{ ...opts, want: BLOCK }` is shorthand for
  // Object.assign({}, opts, { want: BLOCK }) - the same object, spelled shorter.

  if (!allowed || granted === 0) return { allowed: false, source: 'global' };

  leases.set(key, { permits: granted - 1, expiresAt: now + opts.windowMs });
  return { allowed: true, source: 'global' };
}

// Unspent permits must expire with the window, or a process holding a stale
// lease keeps admitting traffic the budget no longer covers.
setInterval(() => {
  const now = Date.now();
  for (const [key, l] of leases) if (l.expiresAt <= now) leases.delete(key);
}, 1000).unref();

// Token bucket in Redis. Unlike a window counter, this one stores no counts:
// it stores the token level and when it was last refilled, and derives the
// current level on read. A client that has been idle may spend its whole
// bucket at once, which is the point - batch jobs arrive in batches.
const TOKEN_BUCKET = `
local capacity = tonumber(ARGV[1])   -- maximum burst
local refill   = tonumber(ARGV[2])   -- tokens added per second
local now      = tonumber(ARGV[3])   -- milliseconds
local cost     = tonumber(ARGV[4])   -- tokens this request costs

local state  = redis.call('HMGET', KEYS[1], 'tokens', 'ts')
local tokens = tonumber(state[1])
local ts     = tonumber(state[2])
if tokens == nil then tokens = capacity; ts = now end

-- Refill lazily: no background job, no timer, just arithmetic on elapsed time.
tokens = math.min(capacity, tokens + ((now - ts) / 1000) * refill)

if tokens < cost then
  local waitMs = math.ceil(((cost - tokens) / refill) * 1000)
  return { 0, 0, waitMs }
end

tokens = tokens - cost
redis.call('HSET', KEYS[1], 'tokens', tokens, 'ts', now)
-- Expire after the time it takes an empty bucket to refill completely; past
-- that point a missing key and a full key mean the same thing.
redis.call('PEXPIRE', KEYS[1], math.ceil((capacity / refill) * 1000) + 1000)
return { 1, math.floor(tokens), 0 }
`;

// A write endpoint costing more than a read is just a different `cost`:
//   consumeTokens(key, { capacity: 100, refill: 10, cost: 1 })   // GET
//   consumeTokens(key, { capacity: 100, refill: 10, cost: 10 })  // bulk import