From 83bfe36c20536f05ffdcb518beacc0ecb63073a2 Mon Sep 17 00:00:00 2001 From: Mohammed Taukir Sheikh Date: Sat, 25 Jul 2026 20:27:44 +0545 Subject: [PATCH] fix(redis): stale-job reaping, single-writer reservoir refresh, no double-execution on release failure Three reviewed fixes plus two pre-existing bugs uncovered by their tests: 1. Stale-job reaping: a crashed process leaked running/currentWeight forever (heartbeats kept the state alive so it never expired). ACQUIRE now stamps job start times into a :jobs:started zset; HEARTBEAT reaps jobs older than the new staleJobTimeout option (default 60s), reclaims counters (clamped at 0), tracks a 'reaped' stat, and the limiter emits a 'reaped' event. RELEASE only decrements when the job still holds its hash entry, so a reaped job's late release cannot drift counters negative. 2. Reservoir refresh multiplied by process count: every instance ran its own setInterval calling UPDATE_RESERVOIR, refilling N times per interval. Refresh is now lazy inside the atomic ACQUIRE script (single writer by construction, keyed off lastReservoirRefresh); the client-side timer is gone. Idle systems refresh on the next acquire. 3. Success-path release failure caused double execution: a Redis error on the success release fell into the catch, released a second time and retried a job whose fn had already succeeded. executeJob now runs fn, releases best-effort exactly once, and settles: success always resolves with the result (release failure only emits 'error'). Found while testing (both pre-existing, both could double-execute or starve jobs): - tryProcess set its re-entrancy guard after awaiting initPromise, so two concurrent callers raced into the loop: one job acquired and executed twice, its sibling dequeued but never executed (promise hung forever). The guard is now claimed synchronously. - PriorityQueue broke FIFO for same-millisecond enqueues (heap ties are order-arbitrary): a later job could starve an earlier equal-priority one behind a depleted reservoir. Added an insertion sequence stamp as final tie-break for a strict total order. Also hardened the post-loop idle check (skip after stop; catch getState rejection) and made the priority/reservoir tests deterministic. Review fixes (cross-model review of this commit): - INIT_STATE stamps lastReservoirRefresh at creation time (and ACQUIRE treats a missing/0 stamp as seed-only), so the first acquire no longer clobbers the configured initial reservoir with reservoirRefreshAmount. Regression test uses reservoir != refreshAmount. - ACQUIRE refuses a jobId still tracked as active (new 'duplicate' reason, treated as a blocked poll): keeps acquire/release accounting symmetric so reused job ids can no longer permanently leak running/currentWeight; the stale entry keeps its start score and is reaped after staleJobTimeout. - tryProcess removes the acquired job by id (removeById) instead of dequeuing the heap root, closing the double-execute/hung-sibling window when the root changes during the acquire round-trip; if the job vanished (cancel/abort) the just-acquired slot is released. - tryProcess loop now catches errors (emit 'error' + re-poll) so a Redis blip during acquire is no longer an unhandled rejection and queued jobs resume once Redis recovers. - Lazy-refresh test strengthened into a real regression test: staggered second instance must not re-grant capacity mid-interval after the first instance consumes the whole refreshed allowance. Co-Authored-By: Claude Fable 5 --- docs/reference/api.mdx | 6 +- src/__tests__/redis-limiter.test.ts | 220 +++++++++++++++++++++++++- src/priority-queue.ts | 9 +- src/redis/distributed-rate-limiter.ts | 128 +++++++++------ src/redis/lua-scripts.ts | 102 ++++++++++-- src/redis/redis-storage.ts | 26 ++- src/types.ts | 9 ++ 7 files changed, 427 insertions(+), 73 deletions(-) diff --git a/docs/reference/api.mdx b/docs/reference/api.mdx index b2d320f..7d2e616 100644 --- a/docs/reference/api.mdx +++ b/docs/reference/api.mdx @@ -27,7 +27,7 @@ description: Main methods and options at a glance. | `updateReservoir(value)` / `incrementReservoir(amount)` | Change reservoir | | `isIdle()` | True if nothing queued and nothing running | -**Events:** `queued`, `executing`, `done`, `failed`, `retry`, `dropped`, `depleted`, `idle`, `error`. +**Events:** `queued`, `executing`, `done`, `failed`, `retry`, `dropped`, `depleted`, `idle`, `error`, `reaped` (distributed only — count of stale jobs reclaimed from dead processes). --- @@ -37,7 +37,9 @@ Same options as RateLimiter, plus: **Required:** `redis` — `{ url }` or `{ host, port }` or `{ client }`, and optional `keyPrefix`, `password`, `db`, etc. -**Extra options:** `pollInterval` (default 50), `heartbeatInterval` (default 30000), `clearOnStart` (default false, useful in tests), `readyTimeout` (default `Infinity` — wait until Redis connects; set a number to fail after that many ms). +**Extra options:** `pollInterval` (default 50), `heartbeatInterval` (default 30000), `staleJobTimeout` (default 60000 — jobs running longer are presumed dead and their slots reclaimed on heartbeat; must exceed your longest expected job duration), `clearOnStart` (default false, useful in tests), `readyTimeout` (default `Infinity` — wait until Redis connects; set a number to fail after that many ms). + +**Note:** in the distributed limiter, `reservoirRefreshInterval` refreshes lazily inside the atomic acquire script (exactly once per interval across all processes, not once per process). A fully idle system does not refresh until the next job tries to acquire. **Methods:** Same as RateLimiter, plus: diff --git a/src/__tests__/redis-limiter.test.ts b/src/__tests__/redis-limiter.test.ts index 8a12d94..97ed093 100644 --- a/src/__tests__/redis-limiter.test.ts +++ b/src/__tests__/redis-limiter.test.ts @@ -252,6 +252,9 @@ async function testReservoir() { executed++; return 3; }); + // Blocked by the empty reservoir; stop() below rejects it — swallow so + // the expected rejection doesn't crash the process + job3Promise.catch(() => {}); // Wait for first two await Promise.all([job1, job2]); @@ -287,6 +290,7 @@ async function testReservoir() { const blocked = limiter1.schedule(async () => { executed++; }); + blocked.catch(() => {}); // rejected by stop() below — expected await sleep(200); assertEqual(executed, 2, 'Only 2 jobs should execute'); @@ -326,11 +330,16 @@ async function testPriority() { const order: string[] = []; - // First job blocks + // First job blocks; wait until it is actually running before queueing + // the others, so they contend on priority rather than racing the start + let firstStarted!: () => void; + const firstStartedP = new Promise((r) => (firstStarted = r)); const first = limiter.schedule({ id: 'first' }, async () => { + firstStarted(); await sleep(50); order.push('first'); }); + await firstStartedP; // Queue with priorities const low = limiter.schedule({ id: 'low', priority: Priority.LOW }, async () => { @@ -455,6 +464,212 @@ async function testCancellation() { }); } +async function testStaleJobReaping() { + console.log('\n💀 Stale Job Reaping Tests'); + + await test('reaps slots held by dead processes', async () => { + const limiter = createLimiter({ + maxConcurrent: 2, + heartbeatInterval: 100, + staleJobTimeout: 250, + }); + await limiter.ready(); + + let reapedCount = 0; + limiter.on('reaped', (count) => { + reapedCount += count; + }); + + // Simulate a crashed process: acquire a slot directly and never release it + const storage = limiter.getStorage(); + const result = await storage.acquireSlot(limiter.id, { + maxConcurrent: 2, + minTime: 0, + maxPerInterval: 999999, + interval: 1000, + weight: 2, + jobId: 'dead-job', + }); + assert(result.allowed, 'Dead job should have acquired a slot'); + + const before = await storage.getState(limiter.id); + assertEqual(before.running, 1, `Expected running=1 before reap, got ${before.running}`); + assertEqual(before.currentWeight, 2, `Expected currentWeight=2 before reap, got ${before.currentWeight}`); + + // Wait past staleJobTimeout plus at least one heartbeat tick + await sleep(600); + + const after = await storage.getState(limiter.id); + assertEqual(after.running, 0, `Expected running reclaimed to 0, got ${after.running}`); + assertEqual(after.currentWeight, 0, `Expected currentWeight reclaimed to 0, got ${after.currentWeight}`); + assertEqual(reapedCount, 1, `Expected reaped event with count 1, got ${reapedCount}`); + + await limiter.stop(); + }); + + await test('does not reap healthy running jobs', async () => { + const limiter = createLimiter({ + maxConcurrent: 1, + heartbeatInterval: 50, + staleJobTimeout: 10000, + }); + await limiter.ready(); + + let reaped = 0; + limiter.on('reaped', (count) => { + reaped += count; + }); + + await limiter.schedule(async () => { + await sleep(200); + }); + + assertEqual(reaped, 0, `Healthy job was reaped ${reaped} times`); + + await limiter.stop(); + }); +} + +async function testLazyReservoirRefresh() { + console.log('\n🔄 Lazy Reservoir Refresh Tests'); + + await test('first acquire honors the configured initial reservoir', async () => { + // reservoir !== reservoirRefreshAmount on purpose: with equal values the + // init-clobber bug (lastReservoirRefresh seeded to 0 making the first + // acquire refresh instantly) is invisible + const limiter = createLimiter({ + reservoir: 1, + reservoirRefreshInterval: 60000, + reservoirRefreshAmount: 5, + }); + await limiter.ready(); + + await limiter.schedule(async () => 1); + + const state = await limiter.getState(); + // Broken behavior: first acquire lazily "refreshes" to refreshAmount(5) + // before decrementing, leaving 4 instead of 0 + assertEqual( + state.reservoir, + 0, + `Expected initial reservoir(1) - 1 acquire = 0, got ${state.reservoir} (init clobbered by refreshAmount?)` + ); + + await limiter.stop(); + }); + + await test('refresh grants capacity exactly once per interval across instances', async () => { + const sharedId = `shared-refresh-${Date.now()}`; + const opts = { + reservoir: 2, + reservoirRefreshInterval: 400, + reservoirRefreshAmount: 2, + }; + + const limiter1 = createLimiter({ id: sharedId, ...opts }); + await limiter1.ready(); + + // Drain the initial reservoir via instance 1 + await limiter1.schedule(async () => 1); + await limiter1.schedule(async () => 2); + + let state = await limiter1.getState(); + assertEqual(state.reservoir, 0, `Reservoir should be drained, got ${state.reservoir}`); + + // Create instance 2 staggered from instance 1, so under per-process-timer + // semantics its refresh clock would fire AFTER instance 1 consumes the + // refreshed allowance below — re-granting capacity mid-interval + await sleep(150); + const limiter2 = createLimiter({ id: sharedId, ...opts, clearOnStart: false }); + await limiter2.ready(); + + // Cross one refresh boundary, then let instance 1 consume the ENTIRE + // refreshed allowance + await sleep(320); + await limiter1.schedule(async () => 3); + await limiter1.schedule(async () => 4); + + state = await limiter1.getState(); + assertEqual( + state.reservoir, + 0, + `Expected refreshAmount(2) fully consumed by instance 1, got ${state.reservoir}` + ); + + // Still within the same interval: instance 2 must NOT be able to acquire. + // Under per-process refresh (the pre-fix design) instance 2's own clock + // refills the reservoir again here and this job executes. + let executed = false; + const blocked = limiter2.schedule(async () => { + executed = true; + }); + blocked.catch(() => {}); // rejected by stop() below — expected + + await sleep(200); + + assertEqual( + executed as boolean, + false, + 'Instance 2 acquired mid-interval — reservoir was refreshed more than once per interval' + ); + state = await limiter2.getState(); + assertEqual( + state.reservoir, + 0, + `Expected reservoir still 0 mid-interval, got ${state.reservoir} (double refresh?)` + ); + + await Promise.all([limiter1.stop(), limiter2.stop()]); + }); +} + +async function testSuccessPathReleaseFailure() { + console.log('\n🧯 Success-Path Release Failure Tests'); + + await test('release failure after success never re-executes the job', async () => { + const limiter = createLimiter({ maxConcurrent: 2, retryCount: 3 }); + await limiter.ready(); + + const storage = limiter.getStorage(); + const originalRelease = storage.releaseSlot.bind(storage); + let releaseCalls = 0; + (storage as any).releaseSlot = async ( + limiterId: string, + jobId: string, + weight: number, + success: boolean + ) => { + releaseCalls++; + if (success) { + throw new Error('Simulated Redis failure on release'); + } + return originalRelease(limiterId, jobId, weight, success); + }; + + let fnRuns = 0; + let errorEvent: Error | null = null; + limiter.on('error', (err) => { + errorEvent = err; + }); + + const result = await limiter.schedule(async () => { + fnRuns++; + return 'success-result'; + }); + + assertEqual(result, 'success-result', 'Job should resolve with its result'); + assertEqual(fnRuns, 1, `fn should run exactly once, ran ${fnRuns} times`); + assertEqual(releaseCalls, 1, `releaseSlot should be called once, got ${releaseCalls}`); + assert(errorEvent !== null, 'Should emit an error event for the release failure'); + assert( + (errorEvent as unknown as Error).message.includes('Simulated Redis failure'), + 'Error event should carry the release failure' + ); + + await limiter.stop(); + }); +} + // Main async function main() { console.log('╔════════════════════════════════════════════╗'); @@ -480,6 +695,9 @@ async function main() { await testPriority(); await testState(); await testCancellation(); + await testStaleJobReaping(); + await testLazyReservoirRefresh(); + await testSuccessPathReleaseFailure(); console.log('\n════════════════════════════════════════'); console.log(`Results: ${passed} passed, ${failed} failed, ${skipped} skipped`); diff --git a/src/priority-queue.ts b/src/priority-queue.ts index 2d9cc75..3b5e46a 100644 --- a/src/priority-queue.ts +++ b/src/priority-queue.ts @@ -6,6 +6,7 @@ import type { Job } from './types.js'; */ export class PriorityQueue { private heap: T[] = []; + private seqCounter = 0; /** * Number of items in the queue @@ -25,6 +26,7 @@ export class PriorityQueue { * Add an item to the queue */ enqueue(item: T): number { + item.seq = this.seqCounter++; this.heap.push(item); this.bubbleUp(this.heap.length - 1); return this.heap.length; @@ -112,7 +114,12 @@ export class PriorityQueue { return a.priority - b.priority; } // Then by queue time (FIFO within same priority) - return a.queuedAt - b.queuedAt; + if (a.queuedAt !== b.queuedAt) { + return a.queuedAt - b.queuedAt; + } + // Date.now() has ms resolution, so same-tick enqueues tie — break by + // insertion order for a strict total order (heap ties are not FIFO) + return (a.seq ?? 0) - (b.seq ?? 0); } /** diff --git a/src/redis/distributed-rate-limiter.ts b/src/redis/distributed-rate-limiter.ts index 1af15ba..7f66109 100644 --- a/src/redis/distributed-rate-limiter.ts +++ b/src/redis/distributed-rate-limiter.ts @@ -53,6 +53,15 @@ export interface DistributedRateLimiterOptions extends RateLimiterOptions { */ heartbeatInterval?: number; + /** + * Jobs running longer than this (ms) are presumed dead — their process + * crashed without releasing the slot — and are reaped on the next + * heartbeat, reclaiming running/currentWeight. Must exceed your longest + * expected job duration. + * @default 60000 + */ + staleJobTimeout?: number; + /** * Whether to clear state on start (for testing) * @default false @@ -107,12 +116,12 @@ export class DistributedRateLimiter extends TypedEventEmitter { private readonly defaultRetryDelay: number | ((attempt: number, error: Error) => number); private readonly pollInterval: number; private readonly heartbeatInterval: number; + private readonly staleJobTimeout: number; - // Reservoir + // Reservoir (refresh happens lazily inside the ACQUIRE_SLOT Lua script) private readonly initialReservoir: number | null; - private readonly reservoirRefreshInterval: number | null; + private readonly reservoirRefreshInterval: number; private readonly reservoirRefreshAmount: number; - private reservoirRefreshTimer: ReturnType | null = null; // Redis storage private readonly storage: RedisStorage; @@ -145,10 +154,12 @@ export class DistributedRateLimiter extends TypedEventEmitter { this.defaultRetryDelay = options.retryDelay ?? 0; this.pollInterval = options.pollInterval ?? 50; this.heartbeatInterval = options.heartbeatInterval ?? 30000; + this.staleJobTimeout = options.staleJobTimeout ?? 60000; - // Reservoir + // Reservoir (0 interval = lazy refresh disabled) this.initialReservoir = options.reservoir ?? null; - this.reservoirRefreshInterval = options.reservoirRefreshInterval ?? null; + this.reservoirRefreshInterval = + this.initialReservoir !== null ? (options.reservoirRefreshInterval ?? 0) : 0; this.reservoirRefreshAmount = options.reservoirRefreshAmount ?? (options.reservoir ?? 0); // Create Redis storage @@ -171,26 +182,18 @@ export class DistributedRateLimiter extends TypedEventEmitter { await this.storage.initialize(this.id, this.initialReservoir); - // Start heartbeat + // Start heartbeat (extends TTLs and reaps stale jobs from dead processes) this.heartbeatTimer = setInterval(async () => { try { - await this.storage.heartbeat(this.id, this.defaultTimeout ?? 60000); + const reaped = await this.storage.heartbeat(this.id, this.staleJobTimeout); + if (reaped > 0) { + this.emit('reaped', reaped); + this.tryProcess(); + } } catch (error) { this.emit('error', error as Error); } }, this.heartbeatInterval); - - // Start reservoir refresh if configured - if (this.reservoirRefreshInterval !== null && this.initialReservoir !== null) { - this.reservoirRefreshTimer = setInterval(async () => { - try { - await this.storage.updateReservoir(this.id, this.reservoirRefreshAmount); - this.tryProcess(); - } catch (error) { - this.emit('error', error as Error); - } - }, this.reservoirRefreshInterval); - } } /** @@ -283,12 +286,14 @@ export class DistributedRateLimiter extends TypedEventEmitter { if (this.processing || this.paused || this.stopped) return; if (this.localQueue.isEmpty) return; - // Ensure initialized - await this.initPromise; - + // Claim the guard synchronously — setting it after an await lets two + // callers race into the loop, double-executing one job and dropping another this.processing = true; try { + // Ensure initialized + await this.initPromise; + while (!this.localQueue.isEmpty && !this.paused && !this.stopped) { const job = this.localQueue.peek(); if (!job) break; @@ -301,11 +306,26 @@ export class DistributedRateLimiter extends TypedEventEmitter { interval: this.interval, weight: job.weight, jobId: job.id, + reservoirRefreshInterval: this.reservoirRefreshInterval, + reservoirRefreshAmount: this.reservoirRefreshAmount, }); if (result.allowed) { - // Remove from local queue and execute - this.localQueue.dequeue(); + // The heap root may have changed during the acquire round-trip (a + // higher-priority schedule(), cancel(), or an abort signal firing) + // — remove the exact job we acquired a slot for, never dequeue() + // the possibly-different current root + const removed = this.localQueue.removeById(job.id); + if (!removed) { + // The peeked job vanished (cancelled/aborted) while we were + // acquiring: give the slot back and keep processing + await this.storage + .releaseSlot(this.id, job.id, job.weight, false) + .catch((e) => + this.emit('error', e instanceof Error ? e : new Error(String(e))) + ); + continue; + } this.executeJob(job); } else { // Wait and retry @@ -317,15 +337,29 @@ export class DistributedRateLimiter extends TypedEventEmitter { await sleep(waitTime); } } + } catch (error) { + // tryProcess is fire-and-forget from schedule()/executeJob()/heartbeat/ + // resume() — a rejection escaping here would be an unhandled promise + // rejection (fatal by default in Node >= 15) and would strand every + // queued job. Surface the error and keep polling instead. + this.emit('error', error instanceof Error ? error : new Error(String(error))); + if (!this.stopped) { + setTimeout(() => this.tryProcess(), this.pollInterval); + } } finally { this.processing = false; } - // Check if idle - if (this.localQueue.isEmpty) { - const state = await this.storage.getState(this.id); - if (state.running === 0) { - this.emit('idle', undefined); + // Check if idle (skip when stopped — the connection may already be + // closed; tryProcess is fire-and-forget so a rejection here is unhandled) + if (this.localQueue.isEmpty && !this.stopped) { + try { + const state = await this.storage.getState(this.id); + if (state.running === 0) { + this.emit('idle', undefined); + } + } catch (error) { + this.emit('error', error as Error); } } } @@ -345,6 +379,9 @@ export class DistributedRateLimiter extends TypedEventEmitter { const waitTime = job.startedAt - job.queuedAt; const startTime = Date.now(); + // Step 1: run the job function; capture the outcome without releasing yet + let result: unknown; + let fnError: Error | null = null; try { // Check if aborted if (job.signal?.aborted) { @@ -352,13 +389,25 @@ export class DistributedRateLimiter extends TypedEventEmitter { } // Execute with optional timeout - let result: unknown; if (job.timeout !== null) { result = await Promise.race([job.fn(), createTimeout(job.timeout, job.id)]); } else { result = await job.fn(); } + } catch (error) { + fnError = error instanceof Error ? error : new Error(String(error)); + } + + // Step 2: release the slot (best-effort; a release failure must never + // retry or reject a job whose fn already succeeded) + try { + await this.storage.releaseSlot(this.id, job.id, job.weight, fnError === null); + } catch (releaseErr) { + this.emit('error', releaseErr instanceof Error ? releaseErr : new Error(String(releaseErr))); + } + // Step 3: settle the job + if (fnError === null) { const duration = Date.now() - startTime; // Update local stats @@ -366,20 +415,10 @@ export class DistributedRateLimiter extends TypedEventEmitter { this.totalExecutionTime += duration; this.localDone++; - // Release slot in Redis - await this.storage.releaseSlot(this.id, job.id, job.weight, true); - this.emit('done', { job, result, duration }); job.resolve(result); - } catch (error) { - const err = error instanceof Error ? error : new Error(String(error)); - - // Release slot in Redis (best-effort; don't leave job promise hanging if Redis fails) - try { - await this.storage.releaseSlot(this.id, job.id, job.weight, false); - } catch (releaseErr) { - this.emit('error', releaseErr instanceof Error ? releaseErr : new Error(String(releaseErr))); - } + } else { + const err = fnError; // Check for retry if (job.retryAttempt < job.retryCount) { @@ -441,11 +480,6 @@ export class DistributedRateLimiter extends TypedEventEmitter { this.heartbeatTimer = null; } - if (this.reservoirRefreshTimer) { - clearInterval(this.reservoirRefreshTimer); - this.reservoirRefreshTimer = null; - } - // Reject all local queued jobs const error = new Error('Limiter was stopped'); while (!this.localQueue.isEmpty) { diff --git a/src/redis/lua-scripts.ts b/src/redis/lua-scripts.ts index adcea73..52dad3c 100644 --- a/src/redis/lua-scripts.ts +++ b/src/redis/lua-scripts.ts @@ -15,6 +15,8 @@ * ARGV[5]: current timestamp (ms) * ARGV[6]: job weight * ARGV[7]: job id + * ARGV[8]: reservoir refresh interval (ms, 0 = disabled) + * ARGV[9]: reservoir refresh amount */ export const ACQUIRE_SLOT = ` local stateKey = KEYS[1] @@ -25,6 +27,8 @@ local interval = tonumber(ARGV[4]) local now = tonumber(ARGV[5]) local weight = tonumber(ARGV[6]) local jobId = ARGV[7] +local refreshInterval = tonumber(ARGV[8] or '0') +local refreshAmount = tonumber(ARGV[9] or '0') -- Get current state local running = tonumber(redis.call('HGET', stateKey, 'running') or '0') @@ -34,6 +38,33 @@ local intervalStart = tonumber(redis.call('HGET', stateKey, 'intervalStart') or local intervalCount = tonumber(redis.call('HGET', stateKey, 'intervalCount') or '0') local reservoir = redis.call('HGET', stateKey, 'reservoir') +-- Lazily refresh the reservoir (single-writer by construction: whichever +-- process crosses the interval boundary first does the reset atomically) +if refreshInterval > 0 then + local lastRefresh = tonumber(redis.call('HGET', stateKey, 'lastReservoirRefresh') or '0') + if lastRefresh == 0 then + -- State hash without a refresh stamp (pre-upgrade data, or legacy init): + -- seed the clock to now WITHOUT touching the reservoir, so the current + -- reservoir value survives its first full interval + redis.call('HSET', stateKey, 'lastReservoirRefresh', now) + elseif now - lastRefresh >= refreshInterval then + redis.call('HSET', stateKey, 'reservoir', refreshAmount) + redis.call('HSET', stateKey, 'lastReservoirRefresh', now) + reservoir = refreshAmount + end +end + +-- Refuse a jobId that is still tracked as active: acquire/release accounting +-- must stay symmetric. A second acquire would HINCRBY running/currentWeight +-- while HSET/ZADD on :jobs merely overwrite the existing member, so the pair +-- of releases would only decrement once — permanently leaking a slot. By +-- returning early the stale entry also keeps its original start score, so +-- the heartbeat reaper reclaims it after staleJobTimeout and a retried job +-- with the same id can then proceed cleanly. +if redis.call('HEXISTS', stateKey .. ':jobs', jobId) == 1 then + return {0, running, 0, 'duplicate'} +end + -- Check concurrency limit if currentWeight + weight > maxConcurrent then return {0, running, 0, 'concurrency'} @@ -74,12 +105,14 @@ if reservoir ~= false then redis.call('HINCRBY', stateKey, 'reservoir', -1) end --- Track active job +-- Track active job (hash: id -> weight, zset: id scored by start time for reaping) redis.call('HSET', stateKey .. ':jobs', jobId, weight) +redis.call('ZADD', stateKey .. ':jobs:started', now, jobId) -- Set TTL on state (cleanup after inactivity) redis.call('EXPIRE', stateKey, 3600) redis.call('EXPIRE', stateKey .. ':jobs', 3600) +redis.call('EXPIRE', stateKey .. ':jobs:started', 3600) return {1, running + 1, 0, 'ok'} `; @@ -97,9 +130,16 @@ local weight = tonumber(ARGV[1]) local jobId = ARGV[2] local success = tonumber(ARGV[3]) --- Decrement running count -local running = redis.call('HINCRBY', stateKey, 'running', -1) -redis.call('HINCRBY', stateKey, 'currentWeight', -weight) +-- Remove from active jobs; only decrement counters if the job still held a +-- slot (it may have already been reaped as stale by HEARTBEAT) +local existed = redis.call('HDEL', stateKey .. ':jobs', jobId) +redis.call('ZREM', stateKey .. ':jobs:started', jobId) + +local running = tonumber(redis.call('HGET', stateKey, 'running') or '0') +if existed == 1 then + running = redis.call('HINCRBY', stateKey, 'running', -1) + redis.call('HINCRBY', stateKey, 'currentWeight', -weight) +end -- Update stats if success == 1 then @@ -108,9 +148,6 @@ else redis.call('HINCRBY', stateKey, 'failed', 1) end --- Remove from active jobs -redis.call('HDEL', stateKey .. ':jobs', jobId) - return running `; @@ -173,10 +210,12 @@ return newValue * Initialize limiter state * KEYS[1]: limiter state key * ARGV[1]: reservoir (or -1 for null) + * ARGV[2]: current timestamp (ms) */ export const INIT_STATE = ` local stateKey = KEYS[1] local reservoir = tonumber(ARGV[1]) +local now = tonumber(ARGV[2]) -- Only initialize if not exists if redis.call('EXISTS', stateKey) == 0 then @@ -187,11 +226,16 @@ if redis.call('EXISTS', stateKey) == 0 then redis.call('HSET', stateKey, 'intervalCount', 0) redis.call('HSET', stateKey, 'done', 0) redis.call('HSET', stateKey, 'failed', 0) - + if reservoir >= 0 then redis.call('HSET', stateKey, 'reservoir', reservoir) + -- Stamp the refresh clock at creation time so the first lazy refresh + -- happens one full interval later. Stamping 0 would make the very first + -- acquire satisfy 'now - lastRefresh >= refreshInterval' and clobber the + -- configured initial reservoir with reservoirRefreshAmount. + redis.call('HSET', stateKey, 'lastReservoirRefresh', now) end - + redis.call('EXPIRE', stateKey, 3600) end @@ -207,28 +251,58 @@ local stateKey = KEYS[1] redis.call('DEL', stateKey) redis.call('DEL', stateKey .. ':jobs') +redis.call('DEL', stateKey .. ':jobs:started') redis.call('DEL', stateKey .. ':queue') return 1 `; /** - * Heartbeat - extend TTL and clean up stale jobs + * Heartbeat - extend TTL and reap stale jobs + * Jobs started more than ARGV[2] ms ago are presumed dead (crashed process + * that never released) and their running/currentWeight is reclaimed. + * Returns: number of jobs reaped + * * KEYS[1]: limiter state key * ARGV[1]: current timestamp - * ARGV[2]: job timeout (ms) + * ARGV[2]: stale job timeout (ms) */ export const HEARTBEAT = ` local stateKey = KEYS[1] local now = tonumber(ARGV[1]) local timeout = tonumber(ARGV[2]) +local jobsKey = stateKey .. ':jobs' +local startedKey = stateKey .. ':jobs:started' -- Extend TTL redis.call('EXPIRE', stateKey, 3600) -redis.call('EXPIRE', stateKey .. ':jobs', 3600) +redis.call('EXPIRE', jobsKey, 3600) +redis.call('EXPIRE', startedKey, 3600) + +-- Reap stale jobs +local reaped = 0 +local stale = redis.call('ZRANGEBYSCORE', startedKey, '-inf', now - timeout) +for _, jobId in ipairs(stale) do + local weight = tonumber(redis.call('HGET', jobsKey, jobId) or '0') + redis.call('HINCRBY', stateKey, 'running', -1) + redis.call('HINCRBY', stateKey, 'currentWeight', -weight) + redis.call('HDEL', jobsKey, jobId) + redis.call('ZREM', startedKey, jobId) + reaped = reaped + 1 +end --- Could add stale job cleanup here if needed +if reaped > 0 then + redis.call('HINCRBY', stateKey, 'reaped', reaped) + + -- Clamp counters at zero (defensive against double-decrement) + if tonumber(redis.call('HGET', stateKey, 'running') or '0') < 0 then + redis.call('HSET', stateKey, 'running', 0) + end + if tonumber(redis.call('HGET', stateKey, 'currentWeight') or '0') < 0 then + redis.call('HSET', stateKey, 'currentWeight', 0) + end +end -return redis.call('HGET', stateKey, 'running') or '0' +return reaped `; diff --git a/src/redis/redis-storage.ts b/src/redis/redis-storage.ts index 137125c..ec8636c 100644 --- a/src/redis/redis-storage.ts +++ b/src/redis/redis-storage.ts @@ -17,7 +17,7 @@ export interface AcquireResult { allowed: boolean; running: number; waitTime: number; - reason: 'ok' | 'concurrency' | 'reservoir' | 'interval' | 'minTime'; + reason: 'ok' | 'concurrency' | 'reservoir' | 'interval' | 'minTime' | 'duplicate'; } /** @@ -185,7 +185,7 @@ export class RedisStorage { // Initialize state const key = this.getKey(limiterId); - await this.execScript('initState', 1, key, reservoir ?? -1); + await this.execScript('initState', 1, key, reservoir ?? -1, Date.now()); this.initialized = true; } @@ -209,6 +209,10 @@ export class RedisStorage { interval: number; weight: number; jobId: string; + /** Lazy reservoir refresh interval in ms (0 = disabled) */ + reservoirRefreshInterval?: number; + /** Value the reservoir is reset to on refresh */ + reservoirRefreshAmount?: number; } ): Promise { const key = this.getKey(limiterId); @@ -224,7 +228,9 @@ export class RedisStorage { options.interval, now, options.weight, - options.jobId + options.jobId, + options.reservoirRefreshInterval ?? 0, + options.reservoirRefreshAmount ?? 0 ) as [number, number, number, string]; return { @@ -319,19 +325,23 @@ export class RedisStorage { } /** - * Send heartbeat to keep state alive + * Send heartbeat to keep state alive and reap stale jobs. + * Jobs started more than `staleJobTimeout` ms ago are presumed dead and + * their slots reclaimed. Returns the number of jobs reaped. */ - async heartbeat(limiterId: string, timeout: number): Promise { + async heartbeat(limiterId: string, staleJobTimeout: number): Promise { const key = this.getKey(limiterId); const now = Date.now(); - return await this.execScript( + const reaped = await this.execScript( 'heartbeat', 1, key, now, - timeout - ) as number; + staleJobTimeout + ); + + return Number(reaped); } /** diff --git a/src/types.ts b/src/types.ts index 52eee1f..96e68b9 100644 --- a/src/types.ts +++ b/src/types.ts @@ -36,6 +36,11 @@ export interface RateLimiterOptions { /** * How often to refill the reservoir (in milliseconds) + * + * In the DistributedRateLimiter the refresh happens lazily inside the + * atomic acquire script (single writer across all processes), so a fully + * idle system does not refresh — the reservoir is refilled the moment the + * next job tries to acquire, which is the only time it matters. * @default null (no automatic refill) */ reservoirRefreshInterval?: number | null; @@ -154,6 +159,8 @@ export interface Job { signal?: AbortSignal; queuedAt: number; startedAt?: number; + /** Insertion-order stamp set by PriorityQueue (FIFO tie-break within same ms) */ + seq?: number; } /** @@ -187,6 +194,8 @@ export type RateLimiterEvents = { idle: void; /** Emitted when the limiter is depleted (reservoir empty) */ depleted: void; + /** Emitted when stale jobs from dead processes are reaped (count reclaimed; distributed limiter only) */ + reaped: number; /** Emitted on any error */ error: Error; };