From 7b4d7f8efde48ad7b60a015884551ead0299d79c Mon Sep 17 00:00:00 2001 From: coo1estguy <70365929+coolestcodeguy@users.noreply.github.com> Date: Tue, 29 Sep 2026 14:33:41 +0530 Subject: [PATCH] fix: restore scheduled queues from postgres and recover consumer groups after redis data loss --- src/queue/index.ts | 15 ++ src/rate-limiter/index.ts | 1 + src/redis/index.ts | 31 +++- src/services/ai/main.ts | 4 +- src/services/delivery/main.ts | 7 +- src/services/scheduler/main.ts | 75 +++++++- tests/chaos/redis-wipe.test.ts | 162 +++++++++++++++++ tests/redis-lua.integration.test.ts | 259 ++++++++++++++++++++++++++++ tests/redis.test.ts | 25 --- tests/scheduler-poll.test.ts | 15 ++ 10 files changed, 557 insertions(+), 37 deletions(-) create mode 100644 tests/chaos/redis-wipe.test.ts create mode 100644 tests/redis-lua.integration.test.ts diff --git a/src/queue/index.ts b/src/queue/index.ts index 7ed0e67..df96033 100644 --- a/src/queue/index.ts +++ b/src/queue/index.ts @@ -395,6 +395,21 @@ export class StreamConsumer { ) { break; // Expected during graceful shutdown } + // The group, or the stream itself, is gone: Redis was flushed, restarted + // without persistence, or failed over to an empty replica. Backing off + // would retry the same read forever, so recreate it and keep reading. + if (err instanceof Error && err.message.startsWith("NOGROUP")) { + try { + await this.ensureGroup(); + this.logger?.warn( + { group: this.group, streams: this.streams }, + "consumer group was missing — recreated it", + ); + continue; + } catch (recreateErr) { + this.logger?.error({ err: recreateErr }, "failed to recreate consumer group"); + } + } this.logger?.error({ err }, "error reading from stream"); await new Promise((resolve) => setTimeout(resolve, retryDelay)); retryDelay = Math.min(retryDelay * 2, 30_000); // Exponential backoff up to 30s diff --git a/src/rate-limiter/index.ts b/src/rate-limiter/index.ts index c5df848..f06eee8 100644 --- a/src/rate-limiter/index.ts +++ b/src/rate-limiter/index.ts @@ -196,6 +196,7 @@ export class UserThrottle { targetTime, memberId, ttlSeconds, + Date.now() - windowMs, ); return { allowed: count <= limit, count, limit }; diff --git a/src/redis/index.ts b/src/redis/index.ts index 8f56da8..099c877 100644 --- a/src/redis/index.ts +++ b/src/redis/index.ts @@ -21,6 +21,7 @@ declare module "ioredis" { targetTime: number | string, memberId: string, ttlSeconds: number | string, + pruneBefore?: number | string, ): Promise; schedulerPoll( key: string, @@ -107,22 +108,36 @@ export const LUA_RENEW_LOCK = ` /** * User-level sliding window notification throttle. * - * A member already in the window was counted by an earlier attempt at the same + * Members are sends scored by when they go out, which for a scheduled send is + * in the future. A send at targetTime is checked against the window ending at + * it, (windowStart, targetTime], so sends scheduled later do not count against + * it. Pruning uses pruneBefore (now minus one window), never windowStart: a + * check for a far-future send must not delete sends that went out today. + * + * A member already present was counted by an earlier attempt at the same * message (a replay after a crash), so it is allowed again without counting * twice. Returns 0 in that case, otherwise the would-be count. + * + * ARGV: windowStart, limit, targetTime, member, ttlSeconds, pruneBefore */ export const LUA_USER_THROTTLE = ` - redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[1]) + local limit = tonumber(ARGV[2]) + local targetTime = tonumber(ARGV[3]) + local ttl = tonumber(ARGV[5]) + redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[6] or ARGV[1]) if redis.call("ZSCORE", KEYS[1], ARGV[4]) then return 0 end - local count = redis.call("ZCARD", KEYS[1]) - if tonumber(count) < tonumber(ARGV[2]) then - redis.call("ZADD", KEYS[1], tonumber(ARGV[3]), ARGV[4]) - redis.call("EXPIRE", KEYS[1], tonumber(ARGV[5])) - return tonumber(count) + 1 + local count = redis.call("ZCOUNT", KEYS[1], "(" .. ARGV[1], targetTime) + if count < limit then + redis.call("ZADD", KEYS[1], targetTime, ARGV[4]) + -- Only ever extend: a send today must not cut short the life of a key + -- that is still holding a send scheduled for next week. + if redis.call("TTL", KEYS[1]) < ttl then + redis.call("EXPIRE", KEYS[1], ttl) + end end - return tonumber(count) + 1 + return count + 1 `; /** Polls scheduler zsets with visibility timeouts. */ diff --git a/src/services/ai/main.ts b/src/services/ai/main.ts index 35e201d..843cc68 100644 --- a/src/services/ai/main.ts +++ b/src/services/ai/main.ts @@ -232,7 +232,9 @@ export class AiWorker extends BaseWorker { if (scheduledAt > now) { await this.db.insert(scheduledPayloads).values({ taskId, - payload: taskPayload, + // The send time rides in the payload so the scheduler can re-queue + // this from Postgres if Redis loses its data. + payload: { ...taskPayload, scheduledAt: pending.scheduledAt }, }); const scheduledEnvelope = buildStreamEvent( diff --git a/src/services/delivery/main.ts b/src/services/delivery/main.ts index 9e37d51..edcb78a 100644 --- a/src/services/delivery/main.ts +++ b/src/services/delivery/main.ts @@ -352,11 +352,14 @@ export class DeliveryWorker extends BaseWorker { return; } + // The retry time rides in the payload so the scheduler can re-queue + // this from Postgres if Redis loses its data. + const retryAt = new Date(Date.now() + retryAfterMs).toISOString(); await this.db .insert(scheduledPayloads) .values({ taskId: task.taskId, - payload: task, + payload: { ...task, scheduledAt: retryAt }, }) .onConflictDoUpdate({ target: scheduledPayloads.taskId, @@ -370,7 +373,7 @@ export class DeliveryWorker extends BaseWorker { projectId: task.projectId, enrichedEventId: task.enrichedEventId, taskId: task.taskId, - scheduledAt: new Date(Date.now() + retryAfterMs).toISOString(), + scheduledAt: retryAt, throttleAttemptCount: task.throttleAttemptCount, }, "delivery", diff --git a/src/services/scheduler/main.ts b/src/services/scheduler/main.ts index 31b36cf..c186957 100644 --- a/src/services/scheduler/main.ts +++ b/src/services/scheduler/main.ts @@ -13,7 +13,7 @@ import { } from "@/index.js"; import { getPriorityBucket, type WorkerOptions } from "@/shared/index.js"; import { scheduledPayloads } from "@/db/schema.js"; -import { inArray } from "drizzle-orm"; +import { asc, gt, inArray } from "drizzle-orm"; import { startHealthReporter } from "@/workers/index.js"; import { createPriorityProducers, @@ -146,8 +146,20 @@ export async function executeSchedulerPoll( const cleanupPipeline = redis.pipeline(); const dbCleanupIds: string[] = []; + // After a rebuild the same task can sit in the ZSET twice (the original + // member carries a trace id Postgres does not keep), so release each task + // once and just clear any second copy. + const released = new Set(); + for (let i = 0; i < parsedTasks.length; i++) { const { taskStr, taskId, traceId } = parsedTasks[i]; + if (released.has(taskId)) { + cleanupPipeline.zrem( + `notif:scheduled:zset:${parseInt(taskId.slice(-1), 16) || 0}`, + taskStr, + ); + continue; + } const payload = payloadMap.get(taskId); if (!payload) { @@ -164,6 +176,7 @@ export async function executeSchedulerPoll( buildStreamEvent("notification.dispatched", dispatchPayload, "scheduler", traceId), ); + released.add(taskId); dbCleanupIds.push(taskId); const shard = parseInt(taskId.slice(-1), 16) || 0; cleanupPipeline.zrem(`notif:scheduled:zset:${shard}`, taskStr); @@ -207,6 +220,63 @@ export async function executeSchedulerPoll( } } +/** + * Present while the scheduler ZSETs are known to be intact. Redis losing its + * data (a flush, a restart without persistence, failover to an empty replica) + * takes this key with it, which is how the poll loop notices. + */ +export const SCHEDULER_SENTINEL_KEY = "notif:scheduler:intact"; + +/** + * Re-queues every scheduled send still waiting in Postgres. Each row keeps its + * payload, and with it the send time, so the ZSETs can be rebuilt exactly; + * a legacy row without a send time is released on the next poll. + * + * Only run when the sentinel is missing: a rebuilt member has no trace id and + * so never matches a surviving original, and the poll only de-duplicates + * copies that happen to be claimed together. + */ +export async function rebuildScheduledQueue( + redis: Redis, + db: any, + logger: Logger, + pageSize = 1000, +): Promise { + let after: string | undefined; + let total = 0; + for (;;) { + const rows: { taskId: string; payload: any }[] = await db + .select({ taskId: scheduledPayloads.taskId, payload: scheduledPayloads.payload }) + .from(scheduledPayloads) + .where(after === undefined ? undefined : gt(scheduledPayloads.taskId, after)) + .orderBy(asc(scheduledPayloads.taskId)) + .limit(pageSize); + if (rows.length === 0) break; + + const pipeline = redis.pipeline(); + for (const { taskId, payload } of rows) { + const sendAt = Date.parse(payload?.scheduledAt ?? ""); + const member = JSON.stringify({ taskId, enrichedEventId: payload?.enrichedEventId }); + const shard = parseInt(taskId.slice(-1), 16) || 0; + pipeline.zadd( + `notif:scheduled:zset:${shard}`, + "NX", + Number.isNaN(sendAt) ? Date.now() : sendAt, + member, + ); + } + await pipeline.exec(); + + total += rows.length; + after = rows[rows.length - 1]!.taskId; + if (rows.length < pageSize) break; + } + // Last, so a crash part-way leaves the sentinel missing and the next poll retries. + await redis.set(SCHEDULER_SENTINEL_KEY, "1"); + if (total > 0) logger.warn({ count: total }, "rebuilt scheduled queue from Postgres"); + return total; +} + export async function startSchedulerWorker() { ({ logger, redis, sql, db } = createWorkerRuntime(config, "scheduler")); ({ consumer, pendingScanner } = createStreamConsumer( @@ -240,6 +310,9 @@ export async function startSchedulerWorker() { const pollLoop = async (): Promise => { if (!isPolling) return; try { + if (!(await redis.native.exists(SCHEDULER_SENTINEL_KEY))) { + await rebuildScheduledQueue(redis.native, db, logger); + } const hasMore = await executeSchedulerPoll(redis.native, outboundProducers, logger, db); if (isPolling) { pollTimeout = setTimeout(() => void pollLoop(), hasMore ? 0 : 5000); diff --git a/tests/chaos/redis-wipe.test.ts b/tests/chaos/redis-wipe.test.ts new file mode 100644 index 0000000..2474c6b --- /dev/null +++ b/tests/chaos/redis-wipe.test.ts @@ -0,0 +1,162 @@ +import { describe, it, expect, beforeAll, afterAll } from "vitest"; +import { PostgreSqlContainer, type StartedPostgreSqlContainer } from "@testcontainers/postgresql"; +import { RedisContainer, type StartedRedisContainer } from "@testcontainers/redis"; +import { randomUUID } from "node:crypto"; +import { Redis } from "ioredis"; +import { NotifkitServer } from "../../src/server.js"; +import { createDatabase, runMigrations } from "../../src/db/index.js"; +import { STREAMS, buildStreamEvent } from "../../src/contracts/index.js"; +import type { Transport, DeliveryResult } from "../../src/transport/index.js"; + +/** + * Redis loses everything mid-run: a FLUSHALL, a failover to an empty replica, + * or a restart without persistence. Postgres survives. + * + * What can come back: the pipeline itself (streams and consumer groups are + * recreated on demand) and anything whose state lives in Postgres — scheduled + * sends keep their payload and send time in `scheduled_payloads`. + * + * What cannot: messages sitting in a stream at the moment of the wipe that no + * stage has written to Postgres yet. Those are gone, so none of these tests + * wipe while messages are in flight. + */ + +class RecordingTransport implements Transport { + readonly channel = "email"; + readonly taskIds: string[] = []; + async send(task: { taskId: string }): Promise { + this.taskIds.push(task.taskId); + return { success: true }; + } +} + +async function waitFor( + what: string, + condition: () => Promise | boolean, + timeoutMs: number, +) { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await condition()) return; + await new Promise((r) => setTimeout(r, 250)); + } + throw new Error(`timed out after ${timeoutMs}ms waiting for: ${what}`); +} + +describe("Redis data loss", () => { + let pg: StartedPostgreSqlContainer; + let redisContainer: StartedRedisContainer; + let redis: Redis; + let sql: any; + let server: NotifkitServer; + const transport = new RecordingTransport(); + const projectId = randomUUID(); + + async function send(count: number, extra: Record = {}) { + for (let i = 0; i < count; i++) { + const event = buildStreamEvent( + "notification.requested", + { + projectId, + target: { type: "user", userId: "wipe-user" }, + channels: ["email"], + templateId: "welcome", + // Critical skips the per-user throttle, which is not what is under test. + priority: "critical", + ...extra, + } as any, + "test", + `trace-${randomUUID()}`, + ); + const full = { ...event, id: randomUUID(), timestamp: new Date().toISOString() }; + await redis.xadd(STREAMS.INBOUND_NORMAL, "*", "data", JSON.stringify(full)); + } + } + + async function scheduledZsetSize(): Promise { + let total = 0; + for (let i = 0; i < 16; i++) total += await redis.zcard(`notif:scheduled:zset:${i}`); + return total; + } + + beforeAll(async () => { + [pg, redisContainer] = await Promise.all([ + new PostgreSqlContainer("postgres:15-alpine").start(), + new RedisContainer("redis:7-alpine").start(), + ]); + + const database = createDatabase({ url: pg.getConnectionUri() }); + sql = database.sql; + await runMigrations(database.db); + + const userId = randomUUID(); + await sql`INSERT INTO projects (id, name) VALUES (${projectId}, 'Wipe')`; + await sql`INSERT INTO users (id, project_id, external_id) VALUES (${userId}, ${projectId}, 'wipe-user')`; + await sql`INSERT INTO user_contacts (id, user_id, channel, target, is_primary) VALUES (${randomUUID()}, ${userId}, 'email', 'wipe@example.com', true)`; + await sql`INSERT INTO templates (project_id, id, channel, topics, content) VALUES (${projectId}, 'welcome', 'email', '{}', '{"subject": "Hi", "text": "Hello"}')`; + + redis = new Redis(redisContainer.getConnectionUrl(), { maxRetriesPerRequest: null }); + + server = new NotifkitServer({ + services: ["enricher", "engine", "scheduler", "delivery"], + redisUrl: redisContainer.getConnectionUrl(), + databaseUrl: pg.getConnectionUri(), + logLevel: "warn", + autoMigrate: false, + providers: [transport], + }); + await server.start(); + + // Baseline: the pipeline works before anything is wiped. + await send(5); + await waitFor("baseline deliveries", () => transport.taskIds.length === 5, 30_000); + }, 180_000); + + afterAll(async () => { + await server?.stop().catch(() => {}); + await sql?.end().catch(() => {}); + redis?.disconnect(); + await Promise.all([pg?.stop(), redisContainer?.stop()]); + }, 60_000); + + it("keeps delivering notifications sent after the wipe", async () => { + const before = transport.taskIds.length; + await redis.flushall(); + + await send(10); + await waitFor("10 post-wipe deliveries", () => transport.taskIds.length >= before + 10, 45_000); + + // Give stragglers a moment, then check nothing went out twice. + await new Promise((r) => setTimeout(r, 2000)); + expect(transport.taskIds.length).toBe(before + 10); + expect(new Set(transport.taskIds).size).toBe(transport.taskIds.length); + }, 90_000); + + it("still delivers a notification that was scheduled before the wipe", async () => { + const before = transport.taskIds.length; + const sendAt = new Date(Date.now() + 8_000); + await send(1, { scheduledAt: sendAt.toISOString() }); + + // Wait until it is fully parked: payload in Postgres, task in the ZSET. + await waitFor( + "scheduled payload stored", + async () => (await sql`SELECT count(*)::int AS n FROM scheduled_payloads`)[0].n === 1, + 30_000, + ); + await waitFor( + "scheduled task queued in Redis", + async () => (await scheduledZsetSize()) === 1, + 30_000, + ); + + await redis.flushall(); + + // Postgres still has the payload and its send time, so this is recoverable. + await waitFor( + "scheduled notification delivered after the wipe", + () => transport.taskIds.length === before + 1, + 45_000, + ); + expect(Date.now()).toBeGreaterThanOrEqual(sendAt.getTime()); + }, 120_000); +}); diff --git a/tests/redis-lua.integration.test.ts b/tests/redis-lua.integration.test.ts new file mode 100644 index 0000000..fa0d638 --- /dev/null +++ b/tests/redis-lua.integration.test.ts @@ -0,0 +1,259 @@ +import { describe, it, expect, beforeAll, afterAll, beforeEach } from "vitest"; +import { RedisContainer, type StartedRedisContainer } from "@testcontainers/redis"; +import { RedisClient } from "@/redis/index.js"; +import { UserThrottle, ApiRateLimiter } from "@/rate-limiter/index.js"; + +/** + * The Lua scripts run against a real Redis. Their whole reason to exist is + * atomicity under concurrent callers, which a mock cannot exercise, so most of + * these fire many calls at once and check the total that got through. + */ + +let container: StartedRedisContainer; +let client: RedisClient; +let redis: RedisClient["native"]; + +beforeAll(async () => { + container = await new RedisContainer("redis:7-alpine").start(); + client = new RedisClient({ url: container.getConnectionUrl() }); + redis = client.native; +}, 120_000); + +afterAll(async () => { + await client?.disconnect(); + await container?.stop(); +}); + +beforeEach(async () => { + await redis.flushall(); +}); + +const HOUR = 3600_000; + +describe("throttleUser (per-user sliding window)", () => { + it("lets exactly `limit` distinct messages through and denies the next", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 3 }); + const results = []; + for (let i = 0; i < 5; i++) { + results.push(await throttle.check("p1", "u1", "normal", { messageId: `m${i}` })); + } + expect(results.map((r) => r.allowed)).toEqual([true, true, true, false, false]); + }); + + it("never over-admits when many messages for one user race", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 10 }); + const results = await Promise.all( + Array.from({ length: 60 }, (_, i) => + throttle.check("p1", "u1", "normal", { messageId: `m${i}` }), + ), + ); + expect(results.filter((r) => r.allowed)).toHaveLength(10); + }); + + it("re-admits a replayed message without spending another slot", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 2 }); + await throttle.check("p1", "u1", "normal", { messageId: "a" }); + await throttle.check("p1", "u1", "normal", { messageId: "b" }); + + // At the limit: a crash replay of "a" must still go out, a new message not. + expect((await throttle.check("p1", "u1", "normal", { messageId: "a" })).allowed).toBe(true); + expect((await throttle.check("p1", "u1", "normal", { messageId: "c" })).allowed).toBe(false); + }); + + it("frees slots once earlier sends fall out of the window", async () => { + const key = "throttle:p1:user:u1"; + const t0 = Date.now(); + for (const m of ["a", "b"]) { + await redis.throttleUser(key, t0 - HOUR, 2, t0, m, 7200); + } + expect(await redis.throttleUser(key, t0 - HOUR, 2, t0, "c", 7200)).toBe(3); // denied + + const later = t0 + HOUR + 1; + expect(await redis.throttleUser(key, later - HOUR, 2, later, "c", 7200)).toBe(1); + }); + + it("does not count sends scheduled for next week against a send right now", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 2 }); + const nextWeek = new Date(Date.now() + 7 * 24 * HOUR).toISOString(); + await throttle.check("p1", "u1", "normal", { messageId: "later-1", scheduledAt: nextWeek }); + await throttle.check("p1", "u1", "normal", { messageId: "later-2", scheduledAt: nextWeek }); + + // Nothing has been sent in the last hour, so this must go out. + expect((await throttle.check("p1", "u1", "normal", { messageId: "now" })).allowed).toBe(true); + }); + + it("does not count a send from this morning against one scheduled for tomorrow", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 1 }); + await throttle.check("p1", "u1", "normal", { messageId: "now" }); + + const tomorrow = new Date(Date.now() + 24 * HOUR).toISOString(); + const result = await throttle.check("p1", "u1", "normal", { + messageId: "tomorrow", + scheduledAt: tomorrow, + }); + expect(result.allowed).toBe(true); + }); + + it("does not forget today's sends when a future-dated send is checked", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 2 }); + await throttle.check("p1", "u1", "normal", { messageId: "a1" }); + await throttle.check("p1", "u1", "normal", { messageId: "a2" }); + + const tomorrow = new Date(Date.now() + 24 * HOUR).toISOString(); + await throttle.check("p1", "u1", "normal", { messageId: "b", scheduledAt: tomorrow }); + + // a1 and a2 went out moments ago, so a third immediate send is over the limit. + expect((await throttle.check("p1", "u1", "normal", { messageId: "c" })).allowed).toBe(false); + }); + + it("does not cut short the life of a key still holding next week's send", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 5 }); + const nextWeek = new Date(Date.now() + 7 * 24 * HOUR).toISOString(); + await throttle.check("p1", "u1", "normal", { messageId: "later", scheduledAt: nextWeek }); + await throttle.check("p1", "u1", "normal", { messageId: "now" }); + + // If the key expired in an hour, next week's send would stop counting. + expect(await redis.ttl("throttle:p1:user:u1")).toBeGreaterThan(6 * 24 * 3600); + }); + + it("keeps each project's count separate for the same user id", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 1 }); + expect((await throttle.check("p1", "u1", "normal", { messageId: "a" })).allowed).toBe(true); + expect((await throttle.check("p2", "u1", "normal", { messageId: "a" })).allowed).toBe(true); + }); + + it("gives the key a TTL, so idle users do not accumulate forever", async () => { + const throttle = new UserThrottle({ redis, maxPerHour: 5 }); + await throttle.check("p1", "u1", "normal", { messageId: "a" }); + const ttl = await redis.ttl("throttle:p1:user:u1"); + expect(ttl).toBeGreaterThan(0); + expect(ttl).toBeLessThanOrEqual(3600); + }); +}); + +describe("leaseApiRateLimit (sliding window token leasing)", () => { + const cur = "{rl:p1}:2"; + const prev = "{rl:p1}:1"; + const window = 60_000; + const startOfBucket2 = 2 * window; + + it("never grants more than the limit in total, however requests interleave", async () => { + const grants = await Promise.all( + Array.from({ length: 100 }, () => + redis.leaseApiRateLimit(cur, prev, startOfBucket2 + 10, window, 100, 7), + ), + ); + expect(grants.reduce((a, b) => a + b, 0)).toBe(100); + expect(grants.every((g) => g >= 0 && g <= 7)).toBe(true); + }); + + it("counts the previous window in proportion to how much of it still overlaps", async () => { + await redis.set(prev, 100); + // Halfway through the current bucket, half the previous count still applies. + const granted = await redis.leaseApiRateLimit( + cur, + prev, + startOfBucket2 + window / 2, + window, + 100, + 1000, + ); + expect(granted).toBe(50); + }); + + it("grants nothing for a zero limit", async () => { + expect(await redis.leaseApiRateLimit(cur, prev, startOfBucket2, window, 0, 5)).toBe(0); + }); + + it("expires the bucket so counters do not live forever", async () => { + await redis.leaseApiRateLimit(cur, prev, startOfBucket2, window, 100, 5); + expect(await redis.ttl(cur)).toBeGreaterThan(0); + }); +}); + +describe("ApiRateLimiter across several processes", () => { + it("admits at most the limit in total when the budget is shared by many instances", async () => { + // An hour-long window keeps the run inside one bucket. + const limiters = Array.from( + { length: 5 }, + () => new ApiRateLimiter({ redis, windowMs: HOUR, keyPrefix: "rl-test" }), + ); + const limit = 60; + const decisions: boolean[] = []; + for (let round = 0; round < 40; round++) { + decisions.push(...(await Promise.all(limiters.map((l) => l.check("proj", limit))))); + } + const admitted = decisions.filter(Boolean).length; + expect(admitted).toBeLessThanOrEqual(limit); + // Leasing strands some tokens in each instance, but most of the budget + // must still be usable or the limiter is effectively a lower limit. + expect(admitted).toBeGreaterThanOrEqual(limit * 0.5); + }); +}); + +describe("schedulerPoll (claim due tasks with a visibility timeout)", () => { + const key = "notif:scheduled:zset:0"; + + it("returns only due tasks, oldest first, up to the limit", async () => { + const now = 1_000_000; + await redis.zadd(key, now - 30, "a", now - 20, "b", now - 10, "c", now + 10, "future"); + + expect(await redis.schedulerPoll(key, now, 2, 60_000)).toEqual(["a", "b"]); + expect(await redis.schedulerPoll(key, now, 10, 60_000)).toEqual(["c"]); + }); + + it("hides a claimed task until its visibility timeout, then hands it out again", async () => { + const now = 1_000_000; + await redis.zadd(key, now - 1, "t"); + + expect(await redis.schedulerPoll(key, now, 10, 60_000)).toEqual(["t"]); + expect(await redis.schedulerPoll(key, now + 59_999, 10, 60_000)).toEqual([]); + // A poller that crashed before removing it must not lose the task. + expect(await redis.schedulerPoll(key, now + 60_000, 10, 60_000)).toEqual(["t"]); + }); + + it("never gives the same task to two pollers racing on one shard", async () => { + const now = 1_000_000; + const tasks = Array.from({ length: 50 }, (_, i) => `task-${i}`); + await redis.zadd(key, ...tasks.flatMap((t, i) => [now - 100 + i, t])); + + const claims = await Promise.all( + Array.from({ length: 20 }, () => redis.schedulerPoll(key, now, 5, 60_000)), + ); + const claimed = claims.flat(); + expect(new Set(claimed).size).toBe(claimed.length); + expect(claimed.sort()).toEqual([...tasks].sort()); + }); +}); + +describe("releaseLock / renewLock (owner-checked lock)", () => { + const key = "lock:test"; + + it("does not let a non-owner release or extend the lock", async () => { + await redis.set(key, "owner-a", "EX", 30); + + expect(await redis.releaseLock(key, "owner-b")).toBe(0); + expect(await redis.renewLock(key, "owner-b", 300)).toBe(0); + expect(await redis.get(key)).toBe("owner-a"); + expect(await redis.ttl(key)).toBeLessThanOrEqual(30); + }); + + it("lets the owner extend and then release it", async () => { + await redis.set(key, "owner-a", "EX", 30); + + expect(await redis.renewLock(key, "owner-a", 300)).toBe(1); + expect(await redis.ttl(key)).toBeGreaterThan(30); + expect(await redis.releaseLock(key, "owner-a")).toBe(1); + expect(await redis.exists(key)).toBe(0); + }); + + it("does not delete a lock that expired and was taken over by someone else", async () => { + await redis.set(key, "owner-a", "PX", 50); + await new Promise((r) => setTimeout(r, 100)); + await redis.set(key, "owner-b", "EX", 30, "NX"); + + // owner-a finishes late and tries to release what it thinks is its lock. + expect(await redis.releaseLock(key, "owner-a")).toBe(0); + expect(await redis.get(key)).toBe("owner-b"); + }); +}); diff --git a/tests/redis.test.ts b/tests/redis.test.ts index dfba3be..8b6058a 100644 --- a/tests/redis.test.ts +++ b/tests/redis.test.ts @@ -72,31 +72,6 @@ describe("Redis Pre-compiled Commands (defineCommand)", () => { void client.disconnect(); } }); - - it("validates Lua scripts contain required Redis commands", () => { - expect(LUA_LEASE_API_RATE_LIMIT).toContain('redis.call("INCRBY"'); - expect(LUA_LEASE_API_RATE_LIMIT).toContain('redis.call("EXPIRE"'); - expect(LUA_LEASE_API_RATE_LIMIT).toContain("math.floor(prevCount * weight + currentCount)"); - - expect(LUA_USER_THROTTLE).toContain('redis.call("ZREMRANGEBYSCORE"'); - expect(LUA_USER_THROTTLE).toContain('redis.call("ZCARD"'); - expect(LUA_USER_THROTTLE).toContain('redis.call("ZADD"'); - - expect(LUA_RELEASE_LOCK).toContain("redis.call('GET'"); - expect(LUA_RELEASE_LOCK).toContain("redis.call('DEL'"); - - expect(LUA_RENEW_LOCK).toContain("redis.call('GET'"); - expect(LUA_RENEW_LOCK).toContain("redis.call('EXPIRE'"); - - expect(LUA_SCHEDULER_POLL).toContain("redis.call('ZRANGE'"); - expect(LUA_SCHEDULER_POLL).toContain("redis.call('ZADD'"); - - expect(LUA_ACQUIRE_LEASE).toContain("redis.call('EXISTS'"); - expect(LUA_ACQUIRE_LEASE).toContain("redis.call('SET'"); - - expect(LUA_MARK_PROCESSED).toContain("redis.call('SET'"); - expect(LUA_MARK_PROCESSED).toContain("redis.call('DEL'"); - }); }); describe("UserThrottle", () => { diff --git a/tests/scheduler-poll.test.ts b/tests/scheduler-poll.test.ts index 2bf27f5..8e857dc 100644 --- a/tests/scheduler-poll.test.ts +++ b/tests/scheduler-poll.test.ts @@ -221,6 +221,21 @@ describe("executeSchedulerPoll", () => { expect(batch[0].metadata.traceId).toBe("trace-abc"); }); + it("releases a task once when a rebuild left a second copy in the same poll", async () => { + // The rebuilt member has no trace id, so it is a different ZSET member. + const rebuilt = JSON.stringify({ taskId: "task-3", enrichedEventId: "evt-1" }); + const h = harness({ shards: shardReplies({ 3: [task("task-3"), rebuilt] }) }); + h.mockDb.queueSelect([{ taskId: "task-3", payload: { priority: "normal" } }]); + + await run(h); + + expect(h.producers.normal.publishBatch).toHaveBeenCalledTimes(1); + expect(h.producers.normal.publishBatch.mock.calls[0]![0]).toHaveLength(1); + // Both copies leave the ZSET, or the second would be claimed again later. + expect(h.cleanupPipeline.zrem).toHaveBeenCalledWith("notif:scheduled:zset:3", task("task-3")); + expect(h.cleanupPipeline.zrem).toHaveBeenCalledWith("notif:scheduled:zset:3", rebuilt); + }); + it("routes by the payload's priority", async () => { const h = harness({ shards: shardReplies({ 0: [task("task-a"), task("task-b")] }) }); h.mockDb.queueSelect([