diff --git a/docs/deployment.md b/docs/deployment.md index 2d56c93..0dc0106 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -18,7 +18,9 @@ with a `Retry-After` header (seconds) and an `X-RateLimit-Remaining` header. The token bucket state is held in **process memory** (`Map`). This is correct and sufficient for a **single Next.js instance** (one server process). Under this deployment the effective limit is exactly the configured 30 req/min per -key. +key. The internal bucket store automatically prunes idle entries (buckets that +have fully refilled and been inactive for at least 1 minute) and enforces an +LRU capacity cap (10,000 buckets by default) to keep process memory bounded. ### Production / multi-instance upgrade path diff --git a/lib/rate-limit.ts b/lib/rate-limit.ts index b1ee664..7f7968c 100644 --- a/lib/rate-limit.ts +++ b/lib/rate-limit.ts @@ -4,7 +4,8 @@ * The token-bucket refill/consume algorithm is decoupled from storage via * the RateLimitStore interface. `InMemoryRateLimitStore` (the default) * keeps bucket state in a per-process Map — sufficient for a single - * Next.js instance. If you run more than one instance behind a load + * Next.js instance. Bounded bucket eviction (idle pruning + LRU capacity cap) + * prevents unbounded memory growth. If you run more than one instance behind a load * balancer, each instance keeps its own counters and the effective limit is * multiplied by the instance count. For multi-instance or serverless (edge) * deployments, inject a shared RateLimitStore backed by Redis or another @@ -43,33 +44,99 @@ interface Bucket { tokens: number /** epoch ms of the last refill */ last: number + /** epoch ms of the last access */ + lastAccess: number +} + +export interface InMemoryRateLimitStoreOptions { + /** Maximum number of buckets allowed before LRU eviction. Defaults to 10,000. */ + maxBuckets?: number + /** Inactivity threshold (ms) after which a fully-refilled bucket is evicted. Defaults to full refill time (maxTokens / refillPerMs). */ + idleTimeoutMs?: number } /** - * Default RateLimitStore: per-process Map, sufficient for a single - * Next.js instance. See docs/deployment.md for the multi-instance upgrade - * path. + * Default RateLimitStore: per-process Map with bounded bucket eviction, + * sufficient for a single Next.js instance. See docs/deployment.md for the + * multi-instance upgrade path. */ export class InMemoryRateLimitStore implements RateLimitStore { private readonly buckets = new Map() + private readonly maxBuckets: number + private readonly idleTimeoutMs: number constructor( private readonly maxTokens: number, private readonly refillPerMs: number, - ) {} + options?: InMemoryRateLimitStoreOptions | number, + idleTimeoutMs?: number, + ) { + if (typeof options === 'number') { + this.maxBuckets = options + this.idleTimeoutMs = idleTimeoutMs ?? Math.ceil(maxTokens / refillPerMs) + } else { + this.maxBuckets = options?.maxBuckets ?? 10_000 + this.idleTimeoutMs = options?.idleTimeoutMs ?? Math.ceil(maxTokens / refillPerMs) + } + } + + /** + * Returns the current number of buckets stored in memory (used for tests and inspection). + */ + getBucketCount(): number { + return this.buckets.size + } + + /** + * Clears all rate limit buckets from memory (used for test isolation). + */ + clear(): void { + this.buckets.clear() + } + + private evictStaleOrOverCapacity(now: number): void { + // 1. Sweep stale/idle buckets that are fully refilled and idle for >= idleTimeoutMs + this.buckets.forEach((bucket, key) => { + const elapsedRefill = now - bucket.last + const currentTokens = Math.min( + this.maxTokens, + bucket.tokens + Math.max(0, elapsedRefill) * this.refillPerMs, + ) + const isFullyRefilled = currentTokens >= this.maxTokens + const isIdle = now - bucket.lastAccess >= this.idleTimeoutMs + + if (isFullyRefilled && isIdle) { + this.buckets.delete(key) + } + }) + + // 2. Enforce hard maxBuckets cap via LRU eviction (front of Map is least recently accessed) + while (this.buckets.size >= this.maxBuckets) { + const oldestKey = this.buckets.keys().next().value + if (oldestKey === undefined) break + this.buckets.delete(oldestKey) + } + } async take(key: string, now: number): Promise<{ tokens: number; consumed: boolean }> { let bucket = this.buckets.get(key) + if (!bucket) { - bucket = { tokens: this.maxTokens, last: now } + this.evictStaleOrOverCapacity(now) + bucket = { tokens: this.maxTokens, last: now, lastAccess: now } this.buckets.set(key, bucket) } else { - // refill based on elapsed time + // Re-insert key to maintain LRU access ordering (most recently accessed key is at the end) + this.buckets.delete(key) + this.buckets.set(key, bucket) + + // Refill based on elapsed time const elapsed = now - bucket.last if (elapsed > 0) { bucket.tokens = Math.min(this.maxTokens, bucket.tokens + elapsed * this.refillPerMs) bucket.last = now } + bucket.lastAccess = now } if (bucket.tokens >= 1) { @@ -88,6 +155,23 @@ const REFILL_PER_MS = MAX_TOKENS / WINDOW_MS // keyed by `${scope}:${id}` — e.g. "ip:1.2.3.4" or "wallet:GA…" const defaultStore = new InMemoryRateLimitStore(MAX_TOKENS, REFILL_PER_MS) +/** + * Resets internal rate limiter state for the default store (used for unit test isolation). + */ +export function resetRateLimitStateForTest(): void { + defaultStore.clear() +} + +/** + * Returns the current bucket count of the store (defaults to defaultStore) for testing. + */ +export function getRateLimitBucketCountForTest(store: RateLimitStore = defaultStore): number { + if (store instanceof InMemoryRateLimitStore) { + return store.getBucketCount() + } + return 0 +} + async function take(store: RateLimitStore, key: string): Promise { const { tokens, consumed } = await store.take(key, Date.now()) if (consumed) { @@ -122,3 +206,4 @@ export async function rateLimitRequest( return ipResult } + diff --git a/test/rate-limit.test.ts b/test/rate-limit.test.ts index 7498543..7f72c4f 100644 --- a/test/rate-limit.test.ts +++ b/test/rate-limit.test.ts @@ -1,6 +1,12 @@ import { describe, test, afterEach } from 'node:test' import * as assert from 'node:assert/strict' -import { rateLimitRequest, InMemoryRateLimitStore, type RateLimitStore } from '../lib/rate-limit' +import { + rateLimitRequest, + InMemoryRateLimitStore, + type RateLimitStore, + resetRateLimitStateForTest, + getRateLimitBucketCountForTest, +} from '../lib/rate-limit' function makeReq(opts: { ip?: string; address?: string } = {}): Request { const headers = new Headers() @@ -11,8 +17,7 @@ function makeReq(opts: { ip?: string; address?: string } = {}): Request { } afterEach(() => { - // clear bucket state between tests by re-importing is not trivial here; - // tests are ordered to avoid cross-contamination (fresh keys per test). + resetRateLimitStateForTest() }) describe('rateLimitRequest', () => { @@ -111,6 +116,78 @@ describe('InMemoryRateLimitStore', () => { }) }) +// =========================================================================== +// InMemoryRateLimitStore eviction and memory bounding +// =========================================================================== + +describe('InMemoryRateLimitStore eviction', () => { + test('evicts idle, fully-refilled buckets when new requests arrive', async () => { + const store = new InMemoryRateLimitStore(5, 5 / 60_000, { maxBuckets: 100, idleTimeoutMs: 1_000 }) + + for (let i = 0; i < 50; i++) { + await store.take(`key-${i}`, 0) + } + assert.equal(store.getBucketCount(), 50) + + // Advance time past idleTimeoutMs (1,000ms) and full refill time (60,000ms) + // Next request triggers eviction sweep of all 50 idle buckets + await store.take('key-new', 60_000) + assert.equal(store.getBucketCount(), 1) + }) + + test('evicted key returning later receives a fresh, correctly-initialized bucket', async () => { + const store = new InMemoryRateLimitStore(5, 5 / 60_000, { maxBuckets: 10, idleTimeoutMs: 1_000 }) + + const r1 = await store.take('k1', 0) + assert.equal(r1.consumed, true) + assert.equal(r1.tokens, 4) + + // Advance time by 60,000ms (idleTimeoutMs = 1,000ms, fully refilled to 5) + // Accessing k2 triggers eviction sweep, removing k1 + await store.take('k2', 60_000) + assert.equal(store.getBucketCount(), 1) + + // k1 returns. Gets fresh bucket with maxTokens (5) - 1 = 4 remaining + const r2 = await store.take('k1', 60_001) + assert.equal(r2.consumed, true) + assert.equal(r2.tokens, 4) + assert.equal(store.getBucketCount(), 2) + }) + + test('enforces maxBuckets cap via LRU eviction when active keys exceed capacity', async () => { + const store = new InMemoryRateLimitStore(5, 5 / 60_000, { maxBuckets: 3, idleTimeoutMs: 1_000_000 }) + + await store.take('k1', 0) + await store.take('k2', 0) + await store.take('k3', 0) + assert.equal(store.getBucketCount(), 3) + + // Access k1 again to refresh its LRU position (LRU order: k2, k3, k1) + await store.take('k1', 10) + + // Taking k4 pushes bucket count over maxBuckets (3), evicting least-recently used key ('k2') + await store.take('k4', 20) + assert.equal(store.getBucketCount(), 3) + + // Accessing k2 now creates a new bucket because it was LRU evicted + const r2 = await store.take('k2', 20) + assert.equal(r2.consumed, true) + assert.equal(r2.tokens, 4) + }) + + test('defaultStore bucket count stays bounded and resetRateLimitStateForTest clears it', async () => { + assert.equal(getRateLimitBucketCountForTest(), 0) + + for (let i = 0; i < 20; i++) { + await rateLimitRequest(makeReq({ ip: `192.168.1.${i}` })) + } + assert.equal(getRateLimitBucketCountForTest(), 20) + + resetRateLimitStateForTest() + assert.equal(getRateLimitBucketCountForTest(), 0) + }) +}) + // =========================================================================== // rateLimitRequest with an injected fake store — pins down key-check order, // short-circuit behavior, and the exact retryAfter/remaining derivation @@ -174,3 +251,4 @@ describe('rateLimitRequest with an injected store', () => { assert.equal(result.retryAfter, 1) }) }) +