import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible'; import type { RateLimitWindow, RateLimitConfig, RateLimitRule } from './types'; // Logger interface for type safety interface Logger { info(message: string, meta?: Record): void; error(message: string, meta?: Record): void; warn(message: string, meta?: Record): void; debug(message: string, meta?: Record): void; } interface RateLimiterEntry { windows: Array<{ limiter: RateLimiterRedis; window: RateLimitWindow; }>; cost: number; } export class QueueRateLimiter { private limiters = new Map(); private rules: RateLimitRule[] = []; private readonly logger: Logger; constructor( private redisClient: ReturnType, logger?: Logger ) { this.logger = logger || console; } /** * Add a rate limit rule */ addRule(rule: RateLimitRule): void { this.rules.push(rule); const key = this.getRuleKey(rule.level, rule.queueName, rule.handler, rule.operation); // Extract limits and cost from config const limits = rule.config.limits || []; const cost = rule.config.cost || 1; // Create rate limiters for each window const windows = limits.map((window, index) => { const limiter = new RateLimiterRedis({ storeClient: this.redisClient, keyPrefix: `rl:${key}:${index}`, points: window.points, duration: window.duration, blockDuration: window.blockDuration || 0, }); return { limiter, window }; }); this.limiters.set(key, { windows, cost }); this.logger.info('Rate limit rule added', { level: rule.level, queueName: rule.queueName, handler: rule.handler, operation: rule.operation, windows: limits.length, cost, }); } /** * Check if a job can be processed based on rate limits * Uses hierarchical precedence: operation > handler > queue > global * The most specific matching rule takes precedence * Returns the longest wait time if multiple windows are hit */ async checkLimit( queueName: string, handler: string, operation: string ): Promise<{ allowed: boolean; retryAfter?: number; remainingPoints?: number; appliedRule?: RateLimitRule; cost?: number; }> { const applicableRule = this.getMostSpecificRule(queueName, handler, operation); if (!applicableRule) { return { allowed: true }; } const key = this.getRuleKey( applicableRule.level, applicableRule.queueName, applicableRule.handler, applicableRule.operation ); const limiterEntry = this.limiters.get(key); if (!limiterEntry || limiterEntry.windows.length === 0) { this.logger.warn('Rate limiter not found for rule', { key, rule: applicableRule }); return { allowed: true }; } const consumerKey = this.getConsumerKey(queueName, handler, operation); const cost = limiterEntry.cost; // Check all windows and collect results const windowResults = await Promise.all( limiterEntry.windows.map(async ({ limiter, window }) => { try { // Try to consume points for this window const result = await limiter.consume(consumerKey, cost); return { allowed: true, remainingPoints: result.remainingPoints, retryAfter: 0, }; } catch (rejRes) { if (rejRes instanceof RateLimiterRes) { return { allowed: false, remainingPoints: rejRes.remainingPoints, retryAfter: rejRes.msBeforeNext, }; } throw rejRes; } }) ); // Find if any window rejected the request const rejectedWindow = windowResults.find(r => !r.allowed); if (rejectedWindow) { // Find the longest wait time among all rejected windows const maxRetryAfter = Math.max( ...windowResults.filter(r => !r.allowed).map(r => r.retryAfter || 0) ); this.logger.warn('Rate limit exceeded', { handler, operation, cost, retryAfter: maxRetryAfter, }); return { allowed: false, retryAfter: maxRetryAfter, remainingPoints: rejectedWindow.remainingPoints, appliedRule: applicableRule, cost, }; } // All windows allowed - return the minimum remaining points const minRemainingPoints = Math.min(...windowResults.map(r => r.remainingPoints || 0)); return { allowed: true, remainingPoints: minRemainingPoints, appliedRule: applicableRule, cost, }; } /** * Get the most specific rule that applies to this job * Precedence: operation > handler > queue > global */ private getMostSpecificRule( queueName: string, handler: string, operation: string ): RateLimitRule | undefined { // 1. Check for operation-specific rule (most specific) let rule = this.rules.find( r => r.level === 'operation' && r.queueName === queueName && r.handler === handler && r.operation === operation ); if (rule) { return rule; } // 2. Check for handler-specific rule rule = this.rules.find( r => r.level === 'handler' && r.queueName === queueName && r.handler === handler ); if (rule) { return rule; } // 3. Check for queue-specific rule rule = this.rules.find(r => r.level === 'queue' && r.queueName === queueName); if (rule) { return rule; } // 4. Check for global rule (least specific) rule = this.rules.find(r => r.level === 'global'); return rule; } /** * Get rule key for storing rate limiter */ private getRuleKey( level: string, queueName?: string, handler?: string, operation?: string ): string { switch (level) { case 'global': return 'global'; case 'queue': return `queue:${queueName}`; case 'handler': return `handler:${queueName}:${handler}`; case 'operation': return `operation:${queueName}:${handler}:${operation}`; default: return level; } } /** * Get consumer key for rate limiting (what gets counted) */ private getConsumerKey(queueName: string, handler: string, operation: string): string { return `${queueName}:${handler}:${operation}`; } /** * Get current rate limit status for a queue/handler/operation */ async getStatus( queueName: string, handler: string, operation: string ): Promise<{ queueName: string; handler: string; operation: string; appliedRule?: RateLimitRule; cost?: number; windows?: Array<{ points: number; duration: number; remaining: number; resetIn: number; }>; }> { const applicableRule = this.getMostSpecificRule(queueName, handler, operation); if (!applicableRule) { return { queueName, handler, operation, }; } const key = this.getRuleKey( applicableRule.level, applicableRule.queueName, applicableRule.handler, applicableRule.operation ); const limiterEntry = this.limiters.get(key); if (!limiterEntry || limiterEntry.windows.length === 0) { return { queueName, handler, operation, appliedRule: applicableRule, }; } try { const consumerKey = this.getConsumerKey(queueName, handler, operation); // Get status for all windows const windows = await Promise.all( limiterEntry.windows.map(async ({ limiter, window }) => { const result = await limiter.get(consumerKey); return { points: window.points, duration: window.duration, remaining: result?.remainingPoints ?? window.points, resetIn: result?.msBeforeNext ?? 0, }; }) ); return { queueName, handler, operation, appliedRule: applicableRule, cost: limiterEntry.cost, windows, }; } catch (error) { this.logger.error('Failed to get rate limit status', { queueName, handler, operation, error, }); return { queueName, handler, operation, appliedRule: applicableRule, cost: limiterEntry.cost, }; } } /** * Reset rate limits for a specific consumer */ async reset(queueName: string, handler?: string, operation?: string): Promise { if (handler && operation) { // Reset specific operation const consumerKey = this.getConsumerKey(queueName, handler, operation); const rule = this.getMostSpecificRule(queueName, handler, operation); if (rule) { const key = this.getRuleKey(rule.level, rule.queueName, rule.handler, rule.operation); const limiterEntry = this.limiters.get(key); if (limiterEntry) { // Reset all windows for this consumer await Promise.all( limiterEntry.windows.map(({ limiter }) => limiter.delete(consumerKey)) ); } } } else { // Reset broader scope - this is more complex with the new hierarchy this.logger.warn('Broad reset not implemented yet', { queueName, handler, operation }); } this.logger.info('Rate limits reset', { queueName, handler, operation }); } /** * Get all configured rate limit rules */ getRules(): RateLimitRule[] { return [...this.rules]; } /** * Remove a rate limit rule */ removeRule(level: string, queueName?: string, handler?: string, operation?: string): boolean { const key = this.getRuleKey(level, queueName, handler, operation); const ruleIndex = this.rules.findIndex( r => r.level === level && (!queueName || r.queueName === queueName) && (!handler || r.handler === handler) && (!operation || r.operation === operation) ); if (ruleIndex >= 0) { this.rules.splice(ruleIndex, 1); this.limiters.delete(key); this.logger.info('Rate limit rule removed', { level, queueName, handler, operation }); return true; } return false; } }