diff --git a/src/lib/ratelimit.js b/src/lib/ratelimit.js new file mode 100644 index 0000000000000000000000000000000000000000..e562afedd957cfb499121ace5e64db14012109f1 --- /dev/null +++ b/src/lib/ratelimit.js @@ -0,0 +1,191 @@ +/** + * @typedef {Object} SlidingWindowRateLimiterOptions + * @property {number} max The maximum amount of tokens that can be consumed in the duration + * @property {number} duration The duration in milliseconds + * @property {number} [cleanupInterval=600000] The interval in milliseconds to clean up expired tokens. Set to 0 to disable. + */ +/** + * @typedef {Object} TokenBucketRateLimiterOptions + * @property {number} capacity The maximum amount of tokens that can be consumed in the duration + * @property {number} refillRate The amount of tokens to refill every refillInterval + * @property {number} refillInterval The interval in milliseconds to refill tokens + */ +/** + * @typedef {Object} Bucket + * @property {number} tokens The amount of tokens in the bucket + * @property {number} lastRefill The timestamp of the last refill + */ +/** + * @typedef {...(SlidingWindowRateLimiter|TokenBucketRateLimiter|AllOrNothingRateLimiter|BurstyRateLimiter)} RateLimiterList + */ + +export class SlidingWindowRateLimiter { + /** + * @param {SlidingWindowRateLimiterOptions} options + */ + constructor(options){ + this.options = Object.assign({max: 5, duration: 1000, cleanupInterval: 600000}, options); + this.ratelimits = new Map(); + this._interval = this.options.cleanupInterval > 0 ? setInterval(() => this._cleanUp(), this.options.cleanupInterval) : null; + } + _cleanUp(){ + for(const key of this.ratelimits.keys()){ + setTimeout(()=>{ + if(this._deleteExpired(key) === 0) this.ratelimits.delete(key); + }, 0); + } + } + _deleteExpired(key){ + const limit = this.ratelimits.get(key); + if(!limit) return 0; + const now = Date.now(); + const filtered = limit.filter(({expires}) => expires > now); + this.ratelimits.set(key, filtered); + return filtered.length; + } + canConsume(key){ + return this._deleteExpired(key) < this.options.max; + } + consume(key){ + if(!this.canConsume(key)) return -1; + const limit = this.ratelimits.get(key) || []; + limit.push({ + expires: Date.now() + this.options.duration + }); + this.ratelimits.set(key, limit); + return this.options.max - limit.length; + } + reset(key){ + return this.ratelimits.delete(key); + } + getRemaining(key){ + return this.options.max - this._deleteExpired(key); + } + getRetryAfter(key){ + if(this._deleteExpired(key) < this.options.max) return 0; + const limit = this.ratelimits.get(key) || []; + return Math.min(...limit.map(({expires})=>expires)) - Date.now(); + } + destory(){ + clearInterval(this._interval); + } +} + +export class TokenBucketRateLimiter { + constructor(options){ + this.options = Object.assign({capacity: 5, refillRate: 1, refillInterval: 1000}, options); + /** + * @type {Map<string, Bucket>} + */ + this.buckets = new Map(); + } + _refill(bucket){ + const now = Date.now(); + const refills = Math.floor((now - bucket.lastRefill) / this.options.refillInterval); + bucket.tokens = Math.min(bucket.tokens + refills * this.options.refillRate, this.options.capacity); + bucket.lastRefill += refills * this.options.refillInterval; + } + canConsume(key){ + const bucket = this.buckets.get(key); + if(!bucket) return true; + this._refill(bucket); + return bucket.tokens > 0; + } + consume(key){ + const bucket = this.buckets.get(key); + if(!bucket){ + this.buckets.set(key, { + tokens: this.options.capacity - 1, + lastRefill: Date.now() + }); + return this.options.capacity - 1; + } + this._refill(bucket); + if(bucket.tokens < 1) return -1; + bucket.tokens--; + return bucket.tokens; + } + reset(key){ + return this.buckets.delete(key); + } + getRemaining(key){ + const bucket = this.buckets.get(key); + if(!bucket) return this.options.capacity; + this._refill(bucket); + return bucket.tokens; + } + getRetryAfter(key){ + if(this.canConsume(key)) return 0; + const bucket = this.buckets.get(key); + return bucket.lastRefill + this.options.refillInterval - Date.now(); + } + destory(){} +} + +/** + * When all of the ratelimiters can consume, it will consume. If one of them can't consume + * then this one can't either. When consuming, internally all ratelimiters will consume. + */ +export class AllOrNothingRateLimiter { + /** + * @param {RateLimiterList} ratelimiters + */ + constructor(...ratelimiters){ + this.ratelimiters = ratelimiters; + } + canConsume(key){ + return this.ratelimiters.every(ratelimiter => ratelimiter.canConsume(key)); + } + consume(key){ + if(!this.canConsume(key)) return -1; + return Math.min(...this.ratelimiters.map(ratelimiter => ratelimiter.consume(key))); + } + reset(key){ + return this.ratelimiters.some(ratelimiter => ratelimiter.reset(key)); + } + getRemaining(key){ + return Math.min(...this.ratelimiters.map(ratelimiter => ratelimiter.getRemaining(key))); + } + getRetryAfter(key){ + return Math.max(...this.ratelimiters.map(ratelimiter => ratelimiter.getRetryAfter(key))); + } + destory(){ + for(const ratelimiter of this.ratelimiters) ratelimiter.destory(); + } +} + +/** + * When one of the ratelimiters can consume, it will consume. If the first one can't consume + * then it will try the second one and so on. If none of them can consume, this one can't either. + * When consuming, internally the first ratelimiter that can consume will consume. + */ +export class BurstyRateLimiter { + /** + * @param {RateLimiterList} ratelimiters + */ + constructor(...ratelimiters){ + this.ratelimiters = ratelimiters; + } + canConsume(key){ + return this.ratelimiters.some(ratelimiter => ratelimiter.canConsume(key)); + } + consume(key){ + for(const ratelimiter of this.ratelimiters){ + let remaining = ratelimiter.consume(key); + if(remaining >= 0) return remaining; + } + return -1; + } + reset(key){ + return this.ratelimiters.some(ratelimiter => ratelimiter.reset(key)); + } + getRemaining(key){ + return this.ratelimiters.map(ratelimiter => ratelimiter.getRemaining(key)).reduce((a, b) => a+b, 0); + } + getRetryAfter(key){ + return Math.min(...this.ratelimiters.map(ratelimiter => ratelimiter.getRetryAfter(key))); + } + destory(){ + for(const ratelimiter of this.ratelimiters) ratelimiter.destory(); + } +}