import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible'; import { getLogger } from '@stock-bot/logger'; import type { RateLimitConfig as BaseRateLimitConfig, RateLimitRule } from './types'; const logger = getLogger('rate-limiter'); // Extend the base config to add rate-limiter specific fields export interface RateLimitConfig extends BaseRateLimitConfig { keyPrefix?: string; } export class QueueRateLimiter { private limiters = new Map(); private rules: RateLimitRule[] = []; constructor(private redisClient: ReturnType) {} /** * 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); const limiter = new RateLimiterRedis({ storeClient: this.redisClient, keyPrefix: `rl:${key}`, points: rule.config.points, duration: rule.config.duration, blockDuration: rule.config.blockDuration || 0, }); this.limiters.set(key, limiter); logger.info('Rate limit rule added', { level: rule.level, queueName: rule.queueName, handler: rule.handler, operation: rule.operation, points: rule.config.points, duration: rule.config.duration, }); } /** * 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 */ async checkLimit(queueName: string, handler: string, operation: string): Promise<{ allowed: boolean; retryAfter?: number; remainingPoints?: number; appliedRule?: RateLimitRule; }> { 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 limiter = this.limiters.get(key); if (!limiter) { logger.warn('Rate limiter not found for rule', { key, rule: applicableRule }); return { allowed: true }; } try { const result = await this.consumePoint(limiter, this.getConsumerKey(queueName, handler, operation)); return { ...result, appliedRule: applicableRule, }; } catch (error) { logger.error('Rate limit check failed', { queueName, handler, operation, error }); // On error, allow the request to proceed return { allowed: true }; } } /** * 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; } /** * Consume a point from the rate limiter */ private async consumePoint( limiter: RateLimiterRedis, key: string ): Promise<{ allowed: boolean; retryAfter?: number; remainingPoints?: number }> { try { const result = await limiter.consume(key); return { allowed: true, remainingPoints: result.remainingPoints, }; } catch (rejRes) { if (rejRes instanceof RateLimiterRes) { logger.warn('Rate limit exceeded', { key, retryAfter: rejRes.msBeforeNext, }); return { allowed: false, retryAfter: rejRes.msBeforeNext, remainingPoints: rejRes.remainingPoints, }; } throw rejRes; } } /** * 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; limit?: { level: string; 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 limiter = this.limiters.get(key); if (!limiter) { return { queueName, handler, operation, appliedRule: applicableRule, }; } try { const consumerKey = this.getConsumerKey(queueName, handler, operation); const result = await limiter.get(consumerKey); const limit = { level: applicableRule.level, points: limiter.points, duration: limiter.duration, remaining: result?.remainingPoints ?? limiter.points, resetIn: result?.msBeforeNext ?? 0, }; return { queueName, handler, operation, appliedRule: applicableRule, limit, }; } catch (error) { logger.error('Failed to get rate limit status', { queueName, handler, operation, error }); return { queueName, handler, operation, appliedRule: applicableRule, }; } } /** * 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 limiter = this.limiters.get(key); if (limiter) { await limiter.delete(consumerKey); } } } else { // Reset broader scope - this is more complex with the new hierarchy logger.warn('Broad reset not implemented yet', { queueName, handler, operation }); } 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); logger.info('Rate limit rule removed', { level, queueName, handler, operation }); return true; } return false; } }