reworked queue lib

libs/queue/src/rate-limiter.ts

@@ -0,0 +1,295 @@
+import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible';
+import { getLogger } from '@stock-bot/logger';
+
+const logger = getLogger('rate-limiter');
+
+export interface RateLimitConfig {
+  points: number;        // Number of requests
+  duration: number;      // Per duration in seconds
+  blockDuration?: number; // Block duration in seconds
+  keyPrefix?: string;
+}
+
+export interface RateLimitRule {
+  level: 'global' | 'handler' | 'operation';
+  handler?: string;
+  operation?: string;
+  config: RateLimitConfig;
+}
+
+export class QueueRateLimiter {
+  private limiters = new Map<string, RateLimiterRedis>();
+  private rules: RateLimitRule[] = [];
+
+  constructor(private redisClient: any) {}
+
+  /**
+   * Add a rate limit rule
+   */
+  addRule(rule: RateLimitRule): void {
+    this.rules.push(rule);
+    
+    const key = this.getRuleKey(rule.level, rule.handler, rule.operation);
+    const limiter = new RateLimiterRedis({
+      storeClient: this.redisClient,
+      keyPrefix: rule.config.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,
+      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
+   */
+  async checkLimit(handler: string, operation: string): Promise<{
+    allowed: boolean;
+    retryAfter?: number;
+    remainingPoints?: number;
+  }> {
+    const limiters = this.getApplicableLimiters(handler, operation);
+    
+    if (limiters.length === 0) {
+      return { allowed: true };
+    }
+
+    try {
+      // Check all applicable rate limiters
+      const results = await Promise.all(
+        limiters.map(({ limiter, key }) => this.consumePoint(limiter, key))
+      );
+
+      // All limiters must allow the request
+      const blocked = results.find(r => !r.allowed);
+      if (blocked) {
+        return blocked;
+      }
+
+      // Return the most restrictive remaining points
+      const minRemainingPoints = Math.min(...results.map(r => r.remainingPoints || Infinity));
+      
+      return {
+        allowed: true,
+        remainingPoints: minRemainingPoints === Infinity ? undefined : minRemainingPoints,
+      };
+    } catch (error) {
+      logger.error('Rate limit check failed', { handler, operation, error });
+      // On error, allow the request to proceed
+      return { allowed: true };
+    }
+  }
+
+  /**
+   * 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 applicable rate limiters for a handler/operation
+   */
+  private getApplicableLimiters(handler: string, operation: string): Array<{ limiter: RateLimiterRedis; key: string }> {
+    const applicable: Array<{ limiter: RateLimiterRedis; key: string }> = [];
+    
+    for (const rule of this.rules) {
+      let applies = false;
+      let consumerKey = '';
+      
+      switch (rule.level) {
+        case 'global':
+          // Global limit applies to all
+          applies = true;
+          consumerKey = 'global';
+          break;
+          
+        case 'handler':
+          // Handler limit applies if handler matches
+          if (rule.handler === handler) {
+            applies = true;
+            consumerKey = handler;
+          }
+          break;
+          
+        case 'operation':
+          // Operation limit applies if both handler and operation match
+          if (rule.handler === handler && rule.operation === operation) {
+            applies = true;
+            consumerKey = `${handler}:${operation}`;
+          }
+          break;
+      }
+      
+      if (applies) {
+        const ruleKey = this.getRuleKey(rule.level, rule.handler, rule.operation);
+        const limiter = this.limiters.get(ruleKey);
+        if (limiter) {
+          applicable.push({ limiter, key: consumerKey });
+        }
+      }
+    }
+    
+    return applicable;
+  }
+
+  /**
+   * Get rule key
+   */
+  private getRuleKey(level: string, handler?: string, operation?: string): string {
+    switch (level) {
+      case 'global':
+        return 'global';
+      case 'handler':
+        return `handler:${handler}`;
+      case 'operation':
+        return `operation:${handler}:${operation}`;
+      default:
+        return level;
+    }
+  }
+
+  /**
+   * Get current rate limit status for a handler/operation
+   */
+  async getStatus(handler: string, operation: string): Promise<{
+    handler: string;
+    operation: string;
+    limits: Array<{
+      level: string;
+      points: number;
+      duration: number;
+      remaining: number;
+      resetIn: number;
+    }>;
+  }> {
+    const applicable = this.getApplicableLimiters(handler, operation);
+    
+    const limits = await Promise.all(
+      applicable.map(async ({ limiter, key }) => {
+        const rule = this.rules.find(r => {
+          const ruleKey = this.getRuleKey(r.level, r.handler, r.operation);
+          return this.limiters.get(ruleKey) === limiter;
+        });
+        
+        try {
+          const result = await limiter.get(key);
+          if (!result) {
+            return {
+              level: rule?.level || 'unknown',
+              points: limiter.points,
+              duration: limiter.duration,
+              remaining: limiter.points,
+              resetIn: 0,
+            };
+          }
+          
+          return {
+            level: rule?.level || 'unknown',
+            points: limiter.points,
+            duration: limiter.duration,
+            remaining: result.remainingPoints,
+            resetIn: result.msBeforeNext,
+          };
+        } catch (error) {
+          return {
+            level: rule?.level || 'unknown',
+            points: limiter.points,
+            duration: limiter.duration,
+            remaining: 0,
+            resetIn: 0,
+          };
+        }
+      })
+    );
+    
+    return {
+      handler,
+      operation,
+      limits,
+    };
+  }
+
+  /**
+   * Reset rate limits for a handler/operation
+   */
+  async reset(handler: string, operation?: string): Promise<void> {
+    const applicable = operation
+      ? this.getApplicableLimiters(handler, operation)
+      : this.rules
+          .filter(r => !handler || r.handler === handler)
+          .map(r => {
+            const key = this.getRuleKey(r.level, r.handler, r.operation);
+            const limiter = this.limiters.get(key);
+            return limiter ? { limiter, key: handler || 'global' } : null;
+          })
+          .filter(Boolean) as Array<{ limiter: RateLimiterRedis; key: string }>;
+    
+    await Promise.all(
+      applicable.map(({ limiter, key }) => limiter.delete(key))
+    );
+    
+    logger.info('Rate limits reset', { handler, operation });
+  }
+
+  /**
+   * Get all configured rate limit rules
+   */
+  getRules(): RateLimitRule[] {
+    return [...this.rules];
+  }
+
+  /**
+   * Remove a rate limit rule
+   */
+  removeRule(level: string, handler?: string, operation?: string): boolean {
+    const key = this.getRuleKey(level, handler, operation);
+    const ruleIndex = this.rules.findIndex(r => 
+      r.level === level && 
+      (!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, handler, operation });
+      return true;
+    }
+    
+    return false;
+  }
+}