platformservice/ts/mta/classes.ratelimiter.ts

import { logger } from '../logger.js';

/**
 * Configuration options for rate limiter
 */
export interface IRateLimitConfig {
  /** Maximum tokens per period */
  maxPerPeriod: number;
  
  /** Time period in milliseconds */
  periodMs: number;
  
  /** Whether to apply per domain/key (vs globally) */
  perKey: boolean;
  
  /** Initial token count (defaults to max) */
  initialTokens?: number;
  
  /** Grace tokens to allow occasional bursts */
  burstTokens?: number;
  
  /** Apply global limit in addition to per-key limits */
  useGlobalLimit?: boolean;
}

/**
 * Token bucket for an individual key
 */
interface TokenBucket {
  /** Current number of tokens */
  tokens: number;
  
  /** Last time tokens were refilled */
  lastRefill: number;
  
  /** Total allowed requests */
  allowed: number;
  
  /** Total denied requests */
  denied: number;
}

/**
 * Rate limiter using token bucket algorithm
 * Provides more sophisticated rate limiting with burst handling
 */
export class RateLimiter {
  /** Rate limit configuration */
  private config: IRateLimitConfig;
  
  /** Token buckets per key */
  private buckets: Map<string, TokenBucket> = new Map();
  
  /** Global bucket for non-keyed rate limiting */
  private globalBucket: TokenBucket;
  
  /**
   * Create a new rate limiter
   * @param config Rate limiter configuration
   */
  constructor(config: IRateLimitConfig) {
    // Set defaults
    this.config = {
      maxPerPeriod: config.maxPerPeriod,
      periodMs: config.periodMs,
      perKey: config.perKey ?? true,
      initialTokens: config.initialTokens ?? config.maxPerPeriod,
      burstTokens: config.burstTokens ?? 0,
      useGlobalLimit: config.useGlobalLimit ?? false
    };
    
    // Initialize global bucket
    this.globalBucket = {
      tokens: this.config.initialTokens,
      lastRefill: Date.now(),
      allowed: 0,
      denied: 0
    };
    
    // Log initialization
    logger.log('info', `Rate limiter initialized: ${this.config.maxPerPeriod} per ${this.config.periodMs}ms${this.config.perKey ? ' per key' : ''}`);
  }
  
  /**
   * Check if a request is allowed under rate limits
   * @param key Key to check rate limit for (e.g. domain, user, IP)
   * @param cost Token cost (defaults to 1)
   * @returns Whether the request is allowed
   */
  public isAllowed(key: string = 'global', cost: number = 1): boolean {
    // If using global bucket directly, just check that
    if (key === 'global' || !this.config.perKey) {
      return this.checkBucket(this.globalBucket, cost);
    }
    
    // Get the key-specific bucket
    const bucket = this.getBucket(key);
    
    // If we also need to check global limit
    if (this.config.useGlobalLimit) {
      // Both key bucket and global bucket must have tokens
      return this.checkBucket(bucket, cost) && this.checkBucket(this.globalBucket, cost);
    } else {
      // Only need to check the key-specific bucket
      return this.checkBucket(bucket, cost);
    }
  }
  
  /**
   * Check if a bucket has enough tokens and consume them
   * @param bucket The token bucket to check
   * @param cost Token cost
   * @returns Whether tokens were consumed
   */
  private checkBucket(bucket: TokenBucket, cost: number): boolean {
    // Refill tokens based on elapsed time
    this.refillBucket(bucket);
    
    // Check if we have enough tokens
    if (bucket.tokens >= cost) {
      // Use tokens
      bucket.tokens -= cost;
      bucket.allowed++;
      return true;
    } else {
      // Rate limit exceeded
      bucket.denied++;
      return false;
    }
  }
  
  /**
   * Consume tokens for a request (if available)
   * @param key Key to consume tokens for
   * @param cost Token cost (defaults to 1)
   * @returns Whether tokens were consumed
   */
  public consume(key: string = 'global', cost: number = 1): boolean {
    const isAllowed = this.isAllowed(key, cost);
    return isAllowed;
  }
  
  /**
   * Get the remaining tokens for a key
   * @param key Key to check
   * @returns Number of remaining tokens
   */
  public getRemainingTokens(key: string = 'global'): number {
    const bucket = this.getBucket(key);
    this.refillBucket(bucket);
    return bucket.tokens;
  }
  
  /**
   * Get stats for a specific key
   * @param key Key to get stats for
   * @returns Rate limit statistics
   */
  public getStats(key: string = 'global'): {
    remaining: number;
    limit: number;
    resetIn: number;
    allowed: number;
    denied: number;
  } {
    const bucket = this.getBucket(key);
    this.refillBucket(bucket);
    
    // Calculate time until next token
    const resetIn = bucket.tokens < this.config.maxPerPeriod ?
      Math.ceil(this.config.periodMs / this.config.maxPerPeriod) :
      0;
    
    return {
      remaining: bucket.tokens,
      limit: this.config.maxPerPeriod,
      resetIn,
      allowed: bucket.allowed,
      denied: bucket.denied
    };
  }
  
  /**
   * Get or create a token bucket for a key
   * @param key The rate limit key
   * @returns Token bucket
   */
  private getBucket(key: string): TokenBucket {
    if (!this.config.perKey || key === 'global') {
      return this.globalBucket;
    }
    
    if (!this.buckets.has(key)) {
      // Create new bucket
      this.buckets.set(key, {
        tokens: this.config.initialTokens,
        lastRefill: Date.now(),
        allowed: 0,
        denied: 0
      });
    }
    
    return this.buckets.get(key);
  }
  
  /**
   * Refill tokens in a bucket based on elapsed time
   * @param bucket Token bucket to refill
   */
  private refillBucket(bucket: TokenBucket): void {
    const now = Date.now();
    const elapsedMs = now - bucket.lastRefill;
    
    // Calculate how many tokens to add
    const rate = this.config.maxPerPeriod / this.config.periodMs;
    const tokensToAdd = elapsedMs * rate;
    
    if (tokensToAdd >= 0.1) { // Allow for partial token refills
      // Add tokens, but don't exceed the normal maximum (without burst)
      // This ensures burst tokens are only used for bursts and don't refill
      const normalMax = this.config.maxPerPeriod;
      bucket.tokens = Math.min(
        // Don't exceed max + burst
        this.config.maxPerPeriod + (this.config.burstTokens || 0), 
        // Don't exceed normal max when refilling
        Math.min(normalMax, bucket.tokens + tokensToAdd)
      );
      
      // Update last refill time
      bucket.lastRefill = now;
    }
  }
  
  /**
   * Reset rate limits for a specific key
   * @param key Key to reset
   */
  public reset(key: string = 'global'): void {
    if (key === 'global' || !this.config.perKey) {
      this.globalBucket.tokens = this.config.initialTokens;
      this.globalBucket.lastRefill = Date.now();
    } else if (this.buckets.has(key)) {
      const bucket = this.buckets.get(key);
      bucket.tokens = this.config.initialTokens;
      bucket.lastRefill = Date.now();
    }
  }
  
  /**
   * Reset all rate limiters
   */
  public resetAll(): void {
    this.globalBucket.tokens = this.config.initialTokens;
    this.globalBucket.lastRefill = Date.now();
    
    for (const bucket of this.buckets.values()) {
      bucket.tokens = this.config.initialTokens;
      bucket.lastRefill = Date.now();
    }
  }
  
  /**
   * Cleanup old buckets to prevent memory leaks
   * @param maxAge Maximum age in milliseconds
   */
  public cleanup(maxAge: number = 24 * 60 * 60 * 1000): void {
    const now = Date.now();
    let removed = 0;
    
    for (const [key, bucket] of this.buckets.entries()) {
      if (now - bucket.lastRefill > maxAge) {
        this.buckets.delete(key);
        removed++;
      }
    }
    
    if (removed > 0) {
      logger.log('debug', `Cleaned up ${removed} stale rate limit buckets`);
    }
  }
}