moved folders around

libs/services/queue/src/rate-limiter.ts

@@ -0,0 +1,294 @@
+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<string, RateLimiterRedis>();
+  private rules: RateLimitRule[] = [];
+
+  constructor(private redisClient: ReturnType<typeof import('./utils').getRedisConnection>) {}
+
+  /**
+   * 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<void> {
+    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;
+  }
+}