elasticsearch/ts/core/plugins/built-in/rate-limit-plugin.ts

/**
 * Rate Limit Plugin
 *
 * Limits request rate to prevent overwhelming Elasticsearch
 */

import { defaultLogger } from '../../observability/logger.js';
import type { Plugin, PluginContext, RateLimitPluginConfig } from '../types.js';

/**
 * Default configuration
 */
const DEFAULT_CONFIG: Required<RateLimitPluginConfig> = {
  maxRequestsPerSecond: 100,
  burstSize: 10,
  waitForSlot: true,
  maxWaitTime: 5000, // 5 seconds
};

/**
 * Token bucket for rate limiting
 */
class TokenBucket {
  private tokens: number;
  private lastRefill: number;

  constructor(
    private maxTokens: number,
    private refillRate: number // tokens per second
  ) {
    this.tokens = maxTokens;
    this.lastRefill = Date.now();
  }

  /**
   * Try to consume a token
   */
  async tryConsume(waitForToken: boolean, maxWaitTime: number): Promise<boolean> {
    this.refill();

    // If we have tokens available, consume one
    if (this.tokens >= 1) {
      this.tokens -= 1;
      return true;
    }

    // If not waiting, reject immediately
    if (!waitForToken) {
      return false;
    }

    // Calculate wait time for next token
    const waitTime = Math.min((1 / this.refillRate) * 1000, maxWaitTime);

    // Wait for token to be available
    await new Promise((resolve) => setTimeout(resolve, waitTime));

    // Try again after waiting
    this.refill();

    if (this.tokens >= 1) {
      this.tokens -= 1;
      return true;
    }

    return false;
  }

  /**
   * Refill tokens based on time elapsed
   */
  private refill(): void {
    const now = Date.now();
    const timePassed = (now - this.lastRefill) / 1000; // seconds
    const tokensToAdd = timePassed * this.refillRate;

    this.tokens = Math.min(this.tokens + tokensToAdd, this.maxTokens);
    this.lastRefill = now;
  }

  /**
   * Get current token count
   */
  getTokens(): number {
    this.refill();
    return this.tokens;
  }

  /**
   * Reset bucket
   */
  reset(): void {
    this.tokens = this.maxTokens;
    this.lastRefill = Date.now();
  }
}

/**
 * Create rate limit plugin
 */
export function createRateLimitPlugin(config: RateLimitPluginConfig = {}): Plugin {
  const pluginConfig = { ...DEFAULT_CONFIG, ...config };
  const logger = defaultLogger;

  let tokenBucket: TokenBucket;
  let rejectedRequests = 0;
  let delayedRequests = 0;
  let totalWaitTime = 0;

  return {
    name: 'rate-limit',
    version: '1.0.0',
    priority: 95, // Execute very late, right before request

    initialize: () => {
      tokenBucket = new TokenBucket(
        pluginConfig.burstSize,
        pluginConfig.maxRequestsPerSecond
      );

      logger.info('Rate limit plugin initialized', {
        maxRequestsPerSecond: pluginConfig.maxRequestsPerSecond,
        burstSize: pluginConfig.burstSize,
        waitForSlot: pluginConfig.waitForSlot,
      });
    },

    beforeRequest: async (context: PluginContext) => {
      const startTime = Date.now();

      // Try to consume a token
      const acquired = await tokenBucket.tryConsume(
        pluginConfig.waitForSlot,
        pluginConfig.maxWaitTime
      );

      if (!acquired) {
        rejectedRequests++;

        logger.warn('Request rate limited', {
          requestId: context.request.requestId,
          rejectedCount: rejectedRequests,
        });

        // Return null to cancel the request
        return null;
      }

      const waitTime = Date.now() - startTime;

      if (waitTime > 100) {
        // Only log if we actually waited
        delayedRequests++;
        totalWaitTime += waitTime;

        logger.debug('Request delayed by rate limiter', {
          requestId: context.request.requestId,
          waitTime,
          availableTokens: tokenBucket.getTokens(),
        });
      }

      return context;
    },
  };
}