diff --git a/src/rate-limiter/index.ts b/src/rate-limiter/index.ts index 12ab7c2..60ba090 100644 --- a/src/rate-limiter/index.ts +++ b/src/rate-limiter/index.ts @@ -1,5 +1,8 @@ import type { Redis } from "@/index.js"; import { LRUCache } from "@/shared/index.js"; +import { LUA_USER_THROTTLE } from "@/redis/index.js"; + +export { LUA_USER_THROTTLE }; import { randomUUID } from "crypto"; @@ -174,17 +177,6 @@ export class UserThrottle { const windowStart = targetTime - windowMs; const memberId = randomUUID(); - const LUA_THROTTLE = ` - redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[1]) - 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 - end - return tonumber(count) + 1 - `; - // The key must outlive the window it is counting. For a future-dated send // that means surviving until targetTime plus one more window, so a task // scheduled for next week still counts against the right bucket. @@ -194,16 +186,19 @@ export class UserThrottle { Math.ceil((targetTime - Date.now()) / 1000) + windowSeconds, ); - const count = (await this.redis.eval( - LUA_THROTTLE, - 1, - key, - windowStart, - limit, - targetTime, - memberId, - ttlSeconds, - )) as number; + const count = + typeof this.redis.throttleUser === "function" + ? await this.redis.throttleUser(key, windowStart, limit, targetTime, memberId, ttlSeconds) + : ((await this.redis.eval( + LUA_USER_THROTTLE, + 1, + key, + windowStart, + limit, + targetTime, + memberId, + ttlSeconds, + )) as number); return { allowed: count <= limit, count, limit }; } diff --git a/src/redis/index.ts b/src/redis/index.ts index 48cf5b6..5cf707e 100644 --- a/src/redis/index.ts +++ b/src/redis/index.ts @@ -1,6 +1,219 @@ import { Redis, type RedisOptions } from "ioredis"; import type { Logger } from "@/index.js"; +declare module "ioredis" { + interface Redis { + checkApiRateLimit( + currentKey: string, + prevKey: string, + now: number | string, + window: number | string, + maxReqs: number | string, + ttl?: number | string, + ): Promise; + checkApiRateLimit1Key( + baseKey: string, + now: number | string, + window: number | string, + maxReqs: number | string, + ttl?: number | string, + ): Promise; + throttleProvider( + key: string, + now: number | string, + windowSeconds: number | string, + limit: number | string, + member: string, + ): Promise<[number, number]>; + releaseLock(lockKey: string, lockToken: string): Promise; + renewLock(lockKey: string, lockToken: string, ttlSeconds: number | string): Promise; + throttleUser( + key: string, + windowStart: number | string, + limit: number | string, + targetTime: number | string, + memberId: string, + ttlSeconds: number | string, + ): Promise; + schedulerPoll( + key: string, + maxScore: number | string, + limit: number | string, + visibilityTimeout: number | string, + ): Promise; + } + + interface ChainableCommander { + schedulerPoll( + key: string, + maxScore: number | string, + limit: number | string, + visibilityTimeout: number | string, + ): this; + } +} + +/** + * Sliding-window counter approximation for API rate limiting. + */ +export const LUA_CHECK_API_RATE_LIMIT = ` + local currentKey = KEYS[1] + local prevKey = KEYS[2] + local now = tonumber(ARGV[1]) + local window = tonumber(ARGV[2]) + local maxReqs = tonumber(ARGV[3]) + local ttl = tonumber(ARGV[4]) or (math.ceil((window * 2) / 1000) + 60) + + if not now or not window or window <= 0 then + return -1 + end + if not maxReqs or maxReqs <= 0 then + return -1 + end + + if not prevKey then + local currentBucket = math.floor(now / window) + local prevBucket = currentBucket - 1 + currentKey = KEYS[1] .. ":" .. currentBucket + prevKey = KEYS[1] .. ":" .. prevBucket + end + + local currentCount = tonumber(redis.call("GET", currentKey) or "0") + local prevCount = tonumber(redis.call("GET", prevKey) or "0") + + local timeIntoCurrent = now % window + local weight = (window - timeIntoCurrent) / window + local estimated = math.floor(prevCount * weight + currentCount) + + if estimated < maxReqs then + local newCount = redis.call("INCR", currentKey) + if newCount == 1 then + redis.call("EXPIRE", currentKey, ttl) + end + return estimated + 1 + end + + return -1 +`; + +/** + * Sliding-window rate limiter for external provider dispatch. + */ +export const LUA_THROTTLE_PROVIDER = ` + local key = KEYS[1] + local now = tonumber(ARGV[1]) + local windowSeconds = tonumber(ARGV[2]) + local limit = tonumber(ARGV[3]) + local member = ARGV[4] + + local clearBefore = now - (windowSeconds * 1000) + + -- Cleanup expired scores + redis.call('ZREMRANGEBYSCORE', key, 0, clearBefore) + + -- Get current count + local count = redis.call('ZCARD', key) + + if count >= limit then + -- Find the oldest score + local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES') + if oldest and oldest[2] then + return {0, tonumber(oldest[2])} + end + return {0, now} + end + + -- Add new request + redis.call('ZADD', key, now, member) + redis.call('EXPIRE', key, windowSeconds * 2) + return {1, 0} +`; + +/** Release a lock only if we still hold it (value matches our token). */ +export const LUA_RELEASE_LOCK = ` + if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) + end + return 0 +`; + +/** Extend a lock's TTL only if we still hold it. */ +export const LUA_RENEW_LOCK = ` + if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('EXPIRE', KEYS[1], ARGV[2]) + end + return 0 +`; + +/** User-level sliding window notification throttle. */ +export const LUA_USER_THROTTLE = ` + redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[1]) + 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 + end + return tonumber(count) + 1 +`; + +/** Polls scheduler zsets with visibility timeouts. */ +export const LUA_SCHEDULER_POLL = ` + local key = KEYS[1] + local maxScore = tonumber(ARGV[1]) + local limit = tonumber(ARGV[2]) + local visibilityTimeout = tonumber(ARGV[3]) or 0 + local tasks = redis.call('ZRANGE', key, 0, maxScore, 'BYSCORE', 'LIMIT', 0, limit) + if #tasks > 0 then + for i, task in ipairs(tasks) do + redis.call('ZADD', key, maxScore + visibilityTimeout, task) + end + end + return tasks +`; + +/** + * Registers pre-compiled custom commands on an ioredis instance. + * Calls defineCommand so ioredis uses EVALSHA rather than re-transmitting + * full script strings over the wire. + */ +export function registerCustomCommands(redis: Redis): void { + redis.defineCommand("checkApiRateLimit", { + numberOfKeys: 2, + lua: LUA_CHECK_API_RATE_LIMIT, + }); + + redis.defineCommand("checkApiRateLimit1Key", { + numberOfKeys: 1, + lua: LUA_CHECK_API_RATE_LIMIT, + }); + + redis.defineCommand("throttleProvider", { + numberOfKeys: 1, + lua: LUA_THROTTLE_PROVIDER, + }); + + redis.defineCommand("releaseLock", { + numberOfKeys: 1, + lua: LUA_RELEASE_LOCK, + }); + + redis.defineCommand("renewLock", { + numberOfKeys: 1, + lua: LUA_RENEW_LOCK, + }); + + redis.defineCommand("throttleUser", { + numberOfKeys: 1, + lua: LUA_USER_THROTTLE, + }); + + redis.defineCommand("schedulerPoll", { + numberOfKeys: 1, + lua: LUA_SCHEDULER_POLL, + }); +} + export interface RedisClientOptions { url: string; name?: string; @@ -25,6 +238,8 @@ export class RedisClient { ...redisOptions, }); + registerCustomCommands(this.native); + this.native.on("connect", () => { this.logger?.info({ url: redactUrl(url) }, "redis connected"); }); diff --git a/src/services/api/main.ts b/src/services/api/main.ts index 99ea597..e86dc8e 100644 --- a/src/services/api/main.ts +++ b/src/services/api/main.ts @@ -34,58 +34,8 @@ export const API_KEY_INVALIDATION_CHANNEL = "apikey.invalidated"; /** Default window length in milliseconds for API rate limiting (1 minute). */ export const API_RATE_LIMIT_WINDOW_MS = 60_000; -/** - * Lua script implementing a sliding-window counter rate limiter. - * - * Approximates request rate across two contiguous time windows without - * maintaining an O(N) sorted set of individual request timestamps. - * - * KEYS[1]: Current window bucket key - * KEYS[2]: Previous window bucket key (optional if computing from KEYS[1]) - * ARGV[1]: Current timestamp in ms - * ARGV[2]: Window size in ms (e.g. 60000) - * ARGV[3]: Maximum requests permitted in window - * ARGV[4]: TTL in seconds for bucket keys (optional, defaults to 2 * window + 60s) - */ -export const LUA_SLIDING_WINDOW_COUNTER = ` - local currentKey = KEYS[1] - local prevKey = KEYS[2] - local now = tonumber(ARGV[1]) - local window = tonumber(ARGV[2]) - local maxReqs = tonumber(ARGV[3]) - local ttl = tonumber(ARGV[4]) or (math.ceil((window * 2) / 1000) + 60) - - if not now or not window or window <= 0 then - return -1 - end - if not maxReqs or maxReqs <= 0 then - return -1 - end - - if not prevKey then - local currentBucket = math.floor(now / window) - local prevBucket = currentBucket - 1 - currentKey = KEYS[1] .. ":" .. currentBucket - prevKey = KEYS[1] .. ":" .. prevBucket - end - - local currentCount = tonumber(redis.call("GET", currentKey) or "0") - local prevCount = tonumber(redis.call("GET", prevKey) or "0") - - local timeIntoCurrent = now % window - local weight = (window - timeIntoCurrent) / window - local estimated = math.floor(prevCount * weight + currentCount) - - if estimated < maxReqs then - local newCount = redis.call("INCR", currentKey) - if newCount == 1 then - redis.call("EXPIRE", currentKey, ttl) - end - return estimated + 1 - end - - return -1 -`; +import { LUA_CHECK_API_RATE_LIMIT } from "@/redis/index.js"; +export const LUA_SLIDING_WINDOW_COUNTER = LUA_CHECK_API_RATE_LIMIT; /** * Evaluates the sliding window counter approximation for given parameters. @@ -587,15 +537,25 @@ export async function startApiServer() { let count = 0; try { - count = (await redis.native.eval( - LUA_SLIDING_WINDOW_COUNTER, - 2, - currentKey, - prevKey, - nowMs, - API_RATE_LIMIT_WINDOW_MS, - projectRateLimitRpm, - )) as number; + if (typeof redis.native.checkApiRateLimit === "function") { + count = await redis.native.checkApiRateLimit( + currentKey, + prevKey, + nowMs, + API_RATE_LIMIT_WINDOW_MS, + projectRateLimitRpm, + ); + } else { + count = (await redis.native.eval( + LUA_SLIDING_WINDOW_COUNTER, + 2, + currentKey, + prevKey, + nowMs, + API_RATE_LIMIT_WINDOW_MS, + projectRateLimitRpm, + )) as number; + } } catch (err) { logger.warn({ err, projectId }, "project rate limit check failed — allowing request"); } diff --git a/src/services/delivery/throttle.ts b/src/services/delivery/throttle.ts index ebbb3b3..57df567 100644 --- a/src/services/delivery/throttle.ts +++ b/src/services/delivery/throttle.ts @@ -1,40 +1,13 @@ import type { Redis } from "ioredis"; +import { LUA_THROTTLE_PROVIDER as LUA_THROTTLE } from "@/redis/index.js"; + +export { LUA_THROTTLE }; export interface ThrottleResult { allowed: boolean; retryAfterMs: number; } -const LUA_THROTTLE = ` - local key = KEYS[1] - local now = tonumber(ARGV[1]) - local windowSeconds = tonumber(ARGV[2]) - local limit = tonumber(ARGV[3]) - local member = ARGV[4] - - local clearBefore = now - (windowSeconds * 1000) - - -- Cleanup expired scores - redis.call('ZREMRANGEBYSCORE', key, 0, clearBefore) - - -- Get current count - local count = redis.call('ZCARD', key) - - if count >= limit then - -- Find the oldest score - local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES') - if oldest and oldest[2] then - return {0, tonumber(oldest[2])} - end - return {0, now} - end - - -- Add new request - redis.call('ZADD', key, now, member) - redis.call('EXPIRE', key, windowSeconds * 2) - return {1, 0} -`; - export async function throttleProvider( redis: Redis, channel: string, @@ -45,15 +18,24 @@ export async function throttleProvider( const now = Date.now(); const zmember = `${now}:${Math.random()}`; - const result = (await redis.eval( - LUA_THROTTLE, - 1, - key, - now.toString(), - config.windowSeconds.toString(), - config.limit.toString(), - zmember, - )) as [number, number]; + const result = + typeof redis.throttleProvider === "function" + ? await redis.throttleProvider( + key, + now.toString(), + config.windowSeconds.toString(), + config.limit.toString(), + zmember, + ) + : ((await redis.eval( + LUA_THROTTLE, + 1, + key, + now.toString(), + config.windowSeconds.toString(), + config.limit.toString(), + zmember, + )) as [number, number]); const allowed = result[0] === 1; const oldestTimestamp = result[1]; diff --git a/src/services/scheduler/main.ts b/src/services/scheduler/main.ts index b313003..79d9582 100644 --- a/src/services/scheduler/main.ts +++ b/src/services/scheduler/main.ts @@ -105,14 +105,12 @@ export async function executeSchedulerPoll( const pipeline = redis.pipeline(); for (let i = 0; i < 16; i++) { - pipeline.eval( - LUA_SCHEDULER_POLL, - 1, - `notif:scheduled:zset:${i}`, - now, - perShardLimit, - visibilityTimeout, - ); + const shardKey = `notif:scheduled:zset:${i}`; + if (typeof pipeline.schedulerPoll === "function") { + pipeline.schedulerPoll(shardKey, now, perShardLimit, visibilityTimeout); + } else { + pipeline.eval(LUA_SCHEDULER_POLL, 1, shardKey, now, perShardLimit, visibilityTimeout); + } } const results = await pipeline.exec(); @@ -212,15 +210,19 @@ export async function executeSchedulerPoll( return false; } finally { if (acquired) { - // Delete the lock only if we still own it - const lua = ` + if (typeof redis.releaseLock === "function") { + await redis.releaseLock(lockKey, lockOwner).catch(() => {}); + } else { + // Delete the lock only if we still own it + const lua = ` if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) else return 0 end `; - await redis.eval(lua, 1, lockKey, lockOwner).catch(() => {}); + await redis.eval(lua, 1, lockKey, lockOwner).catch(() => {}); + } } } } diff --git a/src/services/workflow/main.ts b/src/services/workflow/main.ts index da27a43..0defb8c 100644 --- a/src/services/workflow/main.ts +++ b/src/services/workflow/main.ts @@ -202,11 +202,19 @@ export class WorkflowWorker extends BaseWorker { // Keep the lock alive while the handler runs; without this a handler that // outlives the TTL lets a second resume execute the same steps in parallel. const renewTimer = setInterval(() => { - void this.redisCli - .eval(LUA_RENEW_LOCK, 1, lockKey, lockToken, String(WORKFLOW_LOCK_TTL_SECONDS)) - .catch((err: unknown) => { - this.logger.warn({ err, instanceId }, "failed to renew workflow lock"); - }); + const renewPromise = + typeof this.redisCli.renewLock === "function" + ? this.redisCli.renewLock(lockKey, lockToken, WORKFLOW_LOCK_TTL_SECONDS) + : this.redisCli.eval( + LUA_RENEW_LOCK, + 1, + lockKey, + lockToken, + String(WORKFLOW_LOCK_TTL_SECONDS), + ); + void renewPromise.catch((err: unknown) => { + this.logger.warn({ err, instanceId }, "failed to renew workflow lock"); + }); }, WORKFLOW_LOCK_RENEW_MS); try { @@ -441,7 +449,11 @@ export class WorkflowWorker extends BaseWorker { } finally { clearInterval(renewTimer); // Compare-and-delete: never release a lock a later process re-acquired. - await this.redisCli.eval(LUA_RELEASE_LOCK, 1, lockKey, lockToken); + if (typeof this.redisCli.releaseLock === "function") { + await this.redisCli.releaseLock(lockKey, lockToken); + } else { + await this.redisCli.eval(LUA_RELEASE_LOCK, 1, lockKey, lockToken); + } } await Promise.all(publishPromises); @@ -505,14 +517,23 @@ export async function startWorkflowWorker() { void (async () => { try { const now = Date.now(); - const tasks = (await redis.native.eval( - LUA_SCHEDULER_POLL, - 1, - "notif:workflow:timers", - now, - 100, - WORKFLOW_TIMER_VISIBILITY_MS, - )) as string[]; + const tasks = ( + typeof redis.native.schedulerPoll === "function" + ? await redis.native.schedulerPoll( + "notif:workflow:timers", + now, + 100, + WORKFLOW_TIMER_VISIBILITY_MS, + ) + : await redis.native.eval( + LUA_SCHEDULER_POLL, + 1, + "notif:workflow:timers", + now, + 100, + WORKFLOW_TIMER_VISIBILITY_MS, + ) + ) as string[]; for (const taskStr of tasks) { const task = JSON.parse(taskStr); diff --git a/src/shared/utils.ts b/src/shared/utils.ts index d06cb22..b4ff57b 100644 --- a/src/shared/utils.ts +++ b/src/shared/utils.ts @@ -20,19 +20,7 @@ export function normaliseTarget(target: string): string { return trimmed.includes("@") ? trimmed.toLowerCase() : trimmed; } -export const LUA_SCHEDULER_POLL = ` - local key = KEYS[1] - local maxScore = tonumber(ARGV[1]) - local limit = tonumber(ARGV[2]) - local visibilityTimeout = tonumber(ARGV[3]) or 0 - local tasks = redis.call('ZRANGE', key, 0, maxScore, 'BYSCORE', 'LIMIT', 0, limit) - if #tasks > 0 then - for i, task in ipairs(tasks) do - redis.call('ZADD', key, maxScore + visibilityTimeout, task) - end - end - return tasks -`; +export { LUA_SCHEDULER_POLL } from "@/redis/index.js"; export const LUA_SCHEDULER_CLAIM = ` local payloadKey = KEYS[1] @@ -47,18 +35,4 @@ export const LUA_SCHEDULER_CLAIM = ` return nil end `; -/** Release a lock only if we still hold it (value matches our token). */ -export const LUA_RELEASE_LOCK = ` - if redis.call('GET', KEYS[1]) == ARGV[1] then - return redis.call('DEL', KEYS[1]) - end - return 0 -`; - -/** Extend a lock's TTL only if we still hold it. */ -export const LUA_RENEW_LOCK = ` - if redis.call('GET', KEYS[1]) == ARGV[1] then - return redis.call('EXPIRE', KEYS[1], ARGV[2]) - end - return 0 -`; +export { LUA_RELEASE_LOCK, LUA_RENEW_LOCK } from "@/redis/index.js"; diff --git a/tests/redis.test.ts b/tests/redis.test.ts new file mode 100644 index 0000000..33a891b --- /dev/null +++ b/tests/redis.test.ts @@ -0,0 +1,395 @@ +import { describe, it, expect, vi } from "vitest"; +import { + registerCustomCommands, + RedisClient, + LUA_CHECK_API_RATE_LIMIT, + LUA_THROTTLE_PROVIDER, + LUA_RELEASE_LOCK, + LUA_RENEW_LOCK, + LUA_USER_THROTTLE, + LUA_SCHEDULER_POLL, +} from "@/redis/index.js"; +import { throttleProvider } from "@/services/delivery/throttle.js"; +import { UserThrottle } from "@/rate-limiter/index.js"; +import { executeSchedulerPoll } from "@/services/scheduler/main.js"; +import { WorkflowWorker } from "@/services/workflow/main.js"; + +describe("Redis Pre-compiled Commands (defineCommand)", () => { + describe("Registration & Initialization", () => { + it("registers all custom commands with correct key counts and Lua bodies", () => { + const definedCommands = new Map(); + const mockRedis: any = { + defineCommand: vi.fn((name, def) => { + definedCommands.set(name, def); + }), + }; + + registerCustomCommands(mockRedis); + + expect(definedCommands.has("checkApiRateLimit")).toBe(true); + expect(definedCommands.get("checkApiRateLimit")?.numberOfKeys).toBe(2); + expect(definedCommands.get("checkApiRateLimit")?.lua).toBe(LUA_CHECK_API_RATE_LIMIT); + + expect(definedCommands.has("checkApiRateLimit1Key")).toBe(true); + expect(definedCommands.get("checkApiRateLimit1Key")?.numberOfKeys).toBe(1); + expect(definedCommands.get("checkApiRateLimit1Key")?.lua).toBe(LUA_CHECK_API_RATE_LIMIT); + + expect(definedCommands.has("throttleProvider")).toBe(true); + expect(definedCommands.get("throttleProvider")?.numberOfKeys).toBe(1); + expect(definedCommands.get("throttleProvider")?.lua).toBe(LUA_THROTTLE_PROVIDER); + + expect(definedCommands.has("releaseLock")).toBe(true); + expect(definedCommands.get("releaseLock")?.numberOfKeys).toBe(1); + expect(definedCommands.get("releaseLock")?.lua).toBe(LUA_RELEASE_LOCK); + + expect(definedCommands.has("renewLock")).toBe(true); + expect(definedCommands.get("renewLock")?.numberOfKeys).toBe(1); + expect(definedCommands.get("renewLock")?.lua).toBe(LUA_RENEW_LOCK); + + expect(definedCommands.has("throttleUser")).toBe(true); + expect(definedCommands.get("throttleUser")?.numberOfKeys).toBe(1); + expect(definedCommands.get("throttleUser")?.lua).toBe(LUA_USER_THROTTLE); + + expect(definedCommands.has("schedulerPoll")).toBe(true); + expect(definedCommands.get("schedulerPoll")?.numberOfKeys).toBe(1); + expect(definedCommands.get("schedulerPoll")?.lua).toBe(LUA_SCHEDULER_POLL); + }); + + it("RedisClient automatically registers all custom commands on native client", () => { + const client = new RedisClient({ + url: "redis://localhost:6379", + redisOptions: { lazyConnect: true }, + }); + try { + expect(typeof client.native.checkApiRateLimit).toBe("function"); + expect(typeof client.native.checkApiRateLimit1Key).toBe("function"); + expect(typeof client.native.throttleProvider).toBe("function"); + expect(typeof client.native.releaseLock).toBe("function"); + expect(typeof client.native.renewLock).toBe("function"); + expect(typeof client.native.throttleUser).toBe("function"); + expect(typeof client.native.schedulerPoll).toBe("function"); + } finally { + void client.disconnect(); + } + }); + + it("validates Lua scripts contain required Redis commands", () => { + expect(LUA_CHECK_API_RATE_LIMIT).toContain('redis.call("INCR"'); + expect(LUA_CHECK_API_RATE_LIMIT).toContain('redis.call("EXPIRE"'); + expect(LUA_CHECK_API_RATE_LIMIT).toContain("math.floor(prevCount * weight + currentCount)"); + + expect(LUA_THROTTLE_PROVIDER).toContain("redis.call('ZREMRANGEBYSCORE'"); + expect(LUA_THROTTLE_PROVIDER).toContain("redis.call('ZCARD'"); + expect(LUA_THROTTLE_PROVIDER).toContain("redis.call('ZADD'"); + + 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'"); + }); + }); + + describe("throttleProvider", () => { + it("uses pre-compiled command when available on Redis", async () => { + const mockRedis: any = { + throttleProvider: vi.fn().mockResolvedValue([1, 0]), + eval: vi.fn(), + }; + const mockLogger = { warn: vi.fn(), info: vi.fn() }; + + const result = await throttleProvider( + mockRedis, + "email", + { limit: 10, windowSeconds: 60 }, + mockLogger, + ); + + expect(mockRedis.throttleProvider).toHaveBeenCalled(); + expect(mockRedis.eval).not.toHaveBeenCalled(); + expect(result.allowed).toBe(true); + }); + + it("calculates retryAfterMs when pre-compiled command returns throttled", async () => { + const now = Date.now(); + const mockRedis: any = { + throttleProvider: vi.fn().mockResolvedValue([0, now - 500]), + eval: vi.fn(), + }; + const mockLogger = { warn: vi.fn(), info: vi.fn() }; + + const result = await throttleProvider( + mockRedis, + "sms", + { limit: 5, windowSeconds: 10 }, + mockLogger, + ); + + expect(result.allowed).toBe(false); + expect(result.retryAfterMs).toBeGreaterThanOrEqual(9000); + expect(mockLogger.warn).toHaveBeenCalled(); + }); + + it("falls back to redis.eval when pre-compiled command is absent", async () => { + const mockRedis: any = { + eval: vi.fn().mockResolvedValue([1, 0]), + }; + const mockLogger = { warn: vi.fn(), info: vi.fn() }; + + const result = await throttleProvider( + mockRedis, + "email", + { limit: 10, windowSeconds: 60 }, + mockLogger, + ); + + expect(mockRedis.eval).toHaveBeenCalled(); + expect(result.allowed).toBe(true); + }); + }); + + describe("UserThrottle", () => { + it("uses pre-compiled throttleUser when available", async () => { + const mockRedis: any = { + throttleUser: vi.fn().mockResolvedValue(1), + eval: vi.fn(), + }; + const throttle = new UserThrottle({ redis: mockRedis, maxPerHour: 5 }); + const result = await throttle.check("proj_1", "user_1"); + + expect(mockRedis.throttleUser).toHaveBeenCalled(); + expect(mockRedis.eval).not.toHaveBeenCalled(); + expect(result.allowed).toBe(true); + expect(result.count).toBe(1); + }); + + it("returns allowed false when throttleUser reports limit exceeded", async () => { + const mockRedis: any = { + throttleUser: vi.fn().mockResolvedValue(6), + eval: vi.fn(), + }; + const throttle = new UserThrottle({ redis: mockRedis, maxPerHour: 5 }); + const result = await throttle.check("proj_1", "user_1"); + + expect(mockRedis.throttleUser).toHaveBeenCalled(); + expect(result.allowed).toBe(false); + expect(result.count).toBe(6); + }); + + it("falls back to redis.eval when throttleUser is absent", async () => { + const mockRedis: any = { + eval: vi.fn().mockResolvedValue(1), + }; + const throttle = new UserThrottle({ redis: mockRedis, maxPerHour: 5 }); + const result = await throttle.check("proj_1", "user_1"); + + expect(mockRedis.eval).toHaveBeenCalled(); + expect(result.allowed).toBe(true); + }); + }); + + describe("executeSchedulerPoll (schedulerPoll & releaseLock)", () => { + it("calls pre-compiled schedulerPoll on pipeline and releaseLock on redis", async () => { + const pollPipeline = { + schedulerPoll: vi.fn(), + exec: vi.fn().mockResolvedValue(Array.from({ length: 16 }, () => [])), + }; + const cleanupPipeline = { zrem: vi.fn(), exec: vi.fn().mockResolvedValue([]) }; + + let pipelineCount = 0; + const mockRedis: any = { + set: vi.fn().mockResolvedValue("OK"), + pipeline: vi.fn(() => (pipelineCount++ === 0 ? pollPipeline : cleanupPipeline)), + releaseLock: vi.fn().mockResolvedValue(1), + eval: vi.fn(), + }; + const mockProducers = { + critical: { publishBatch: vi.fn() }, + normal: { publishBatch: vi.fn() }, + low: { publishBatch: vi.fn() }, + }; + const mockLogger = { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }; + const mockDb: any = {}; + + const polled = await executeSchedulerPoll(mockRedis, mockProducers, mockLogger, mockDb); + + expect(polled).toBe(false); + expect(pollPipeline.schedulerPoll).toHaveBeenCalledTimes(16); + expect(mockRedis.releaseLock).toHaveBeenCalledWith( + "notif:lock:scheduler:poll", + expect.any(String), + ); + expect(mockRedis.eval).not.toHaveBeenCalled(); + }); + + it("falls back to pipeline.eval and redis.eval when pre-compiled methods are absent", async () => { + const pollPipeline = { + eval: vi.fn(), + exec: vi.fn().mockResolvedValue(Array.from({ length: 16 }, () => [null, []])), + }; + const cleanupPipeline = { zrem: vi.fn(), exec: vi.fn().mockResolvedValue([]) }; + + let pipelineCount = 0; + const mockRedis: any = { + set: vi.fn().mockResolvedValue("OK"), + pipeline: vi.fn(() => (pipelineCount++ === 0 ? pollPipeline : cleanupPipeline)), + eval: vi.fn().mockResolvedValue(1), + }; + const mockProducers = { + critical: { publishBatch: vi.fn() }, + normal: { publishBatch: vi.fn() }, + low: { publishBatch: vi.fn() }, + }; + const mockLogger = { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }; + const mockDb: any = {}; + + const polled = await executeSchedulerPoll(mockRedis, mockProducers, mockLogger, mockDb); + + expect(polled).toBe(false); + expect(pollPipeline.eval).toHaveBeenCalledTimes(16); + expect(mockRedis.eval).toHaveBeenCalled(); + }); + }); + + describe("WorkflowWorker (releaseLock & renewLock)", () => { + it("calls pre-compiled releaseLock when available", async () => { + const { workflowRegistry } = await import("@/workflows/index.js"); + workflowRegistry.register("flow", async () => {}); + + const mockRedis: any = { + set: vi.fn().mockResolvedValue("OK"), + releaseLock: vi.fn().mockResolvedValue(1), + renewLock: vi.fn().mockResolvedValue(1), + eval: vi.fn(), + zadd: vi.fn().mockResolvedValue(1), + }; + + const qb: any = {}; + qb.select = vi.fn().mockReturnValue(qb); + qb.from = vi.fn().mockReturnValue(qb); + qb.where = vi.fn().mockReturnValue(qb); + qb.limit = vi + .fn() + .mockResolvedValue([{ id: "inst-1", status: "pending", name: "flow", input: {} }]); + qb.update = vi.fn().mockReturnValue(qb); + qb.set = vi.fn().mockReturnValue(qb); + qb.then = function (resolve: any) { + resolve([{ id: "inst-1", status: "pending", name: "flow", input: {} }]); + }; + + const worker = new WorkflowWorker({ + consumer: { ack: vi.fn(), nack: vi.fn() } as any, + pendingScanner: {} as any, + logger: { + child: vi.fn().mockReturnThis(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + } as any, + concurrency: 1, + redis: mockRedis, + db: qb, + workflowProducer: { + publishBatch: vi.fn().mockResolvedValue({ messageIds: [], eventIds: [] }), + }, + notificationProducer: { + publishBatch: vi.fn().mockResolvedValue({ messageIds: [], eventIds: [] }), + }, + }); + + const msg = { + id: "msg-1", + stream: "workflow" as any, + event: { + id: "evt-1", + type: "workflow.triggered", + name: "workflow.triggered", + payload: { instanceId: "inst-1", name: "flow", projectId: "proj-1" }, + metadata: { timestamp: new Date().toISOString() }, + }, + }; + + await worker.process(msg as any); + + expect(mockRedis.releaseLock).toHaveBeenCalledWith( + "lock:workflow:inst-1", + expect.any(String), + ); + expect(mockRedis.eval).not.toHaveBeenCalled(); + }); + + it("falls back to eval for releaseLock when releaseLock is absent", async () => { + const { workflowRegistry } = await import("@/workflows/index.js"); + workflowRegistry.register("flow2", async () => {}); + + const mockRedis: any = { + set: vi.fn().mockResolvedValue("OK"), + eval: vi.fn().mockResolvedValue(1), + zadd: vi.fn().mockResolvedValue(1), + }; + + const qb: any = {}; + qb.select = vi.fn().mockReturnValue(qb); + qb.from = vi.fn().mockReturnValue(qb); + qb.where = vi.fn().mockReturnValue(qb); + qb.limit = vi + .fn() + .mockResolvedValue([{ id: "inst-2", status: "pending", name: "flow2", input: {} }]); + qb.update = vi.fn().mockReturnValue(qb); + qb.set = vi.fn().mockReturnValue(qb); + qb.then = function (resolve: any) { + resolve([{ id: "inst-2", status: "pending", name: "flow2", input: {} }]); + }; + + const worker = new WorkflowWorker({ + consumer: { ack: vi.fn(), nack: vi.fn() } as any, + pendingScanner: {} as any, + logger: { + child: vi.fn().mockReturnThis(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + } as any, + concurrency: 1, + redis: mockRedis, + db: qb, + workflowProducer: { + publishBatch: vi.fn().mockResolvedValue({ messageIds: [], eventIds: [] }), + }, + notificationProducer: { + publishBatch: vi.fn().mockResolvedValue({ messageIds: [], eventIds: [] }), + }, + }); + + const msg = { + id: "msg-2", + stream: "workflow" as any, + event: { + id: "evt-2", + type: "workflow.triggered", + name: "workflow.triggered", + payload: { instanceId: "inst-2", name: "flow2", projectId: "proj-1" }, + metadata: { timestamp: new Date().toISOString() }, + }, + }; + + await worker.process(msg as any); + + expect(mockRedis.eval).toHaveBeenCalledWith( + expect.stringContaining("redis.call('GET'"), + 1, + "lock:workflow:inst-2", + expect.any(String), + ); + }); + }); +});