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

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;
  }
}