/** * 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 = { 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 { 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; }, }; }