update
This commit is contained in:
281
ts/mail/delivery/classes.ratelimiter.ts
Normal file
281
ts/mail/delivery/classes.ratelimiter.ts
Normal file
@ -0,0 +1,281 @@
|
||||
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`);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user