stock-bot/libs/core/queue/src/rate-limiter.ts

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<string, unknown>): void;
  error(message: string, meta?: Record<string, unknown>): void;
  warn(message: string, meta?: Record<string, unknown>): void;
  debug(message: string, meta?: Record<string, unknown>): void;
}

interface RateLimiterEntry {
  windows: Array<{
    limiter: RateLimiterRedis;
    window: RateLimitWindow;
  }>;
  cost: number;
}

export class QueueRateLimiter {
  private limiters = new Map<string, RateLimiterEntry>();
  private rules: RateLimitRule[] = [];
  private readonly logger: Logger;

  constructor(
    private redisClient: ReturnType<typeof import('./utils').getRedisConnection>,
    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<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 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;
  }
}