limiter
Version:
A generic rate limiter for the web and node.js. Useful for API clients, web crawling, or other tasks that need to be throttled
114 lines • 5.05 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.RateLimiter = void 0;
const TokenBucket_js_1 = require("./TokenBucket.js");
const clock_js_1 = require("./clock.js");
/**
* A generic rate limiter. Underneath the hood, this uses a token bucket plus
* an additional check to limit how many tokens we can remove each interval.
*
* @param options
* @param options.tokensPerInterval Maximum number of tokens that can be
* removed at any given moment and over the course of one interval.
* @param options.interval The interval length in milliseconds, or as
* one of the following strings: 'second', 'minute', 'hour', day'.
* @param options.fireImmediately Whether or not the promise will resolve
* immediately when rate limiting is in effect (default is false).
*/
class RateLimiter {
tokenBucket;
curIntervalStart;
tokensThisInterval;
fireImmediately;
constructor({ tokensPerInterval, interval, fireImmediately }) {
this.tokenBucket = new TokenBucket_js_1.TokenBucket({
bucketSize: tokensPerInterval,
tokensPerInterval,
interval,
});
// Fill the token bucket to start
this.tokenBucket.content = tokensPerInterval;
this.curIntervalStart = (0, clock_js_1.getMilliseconds)();
this.tokensThisInterval = 0;
this.fireImmediately = fireImmediately ?? false;
}
/**
* Remove the requested number of tokens. If the rate limiter contains enough
* tokens and we haven't spent too many tokens in this interval already, this
* will happen immediately. Otherwise, the removal will happen when enough
* tokens become available.
* @param count The number of tokens to remove.
* @returns A promise for the remainingTokens count.
*/
async removeTokens(count) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize) {
throw new Error(`Requested tokens ${count} exceeds maximum tokens per interval ${this.tokenBucket.bucketSize}`);
}
const now = (0, clock_js_1.getMilliseconds)();
// Advance the current interval and reset the current interval token count
// if needed
if (now < this.curIntervalStart || now - this.curIntervalStart >= this.tokenBucket.interval) {
this.curIntervalStart = now;
this.tokensThisInterval = 0;
}
// If we don't have enough tokens left in this interval, wait until the
// next interval
if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval) {
if (this.fireImmediately) {
return -1;
}
else {
const waitMs = Math.ceil(this.curIntervalStart + this.tokenBucket.interval - now);
await (0, clock_js_1.wait)(waitMs);
const remainingTokens = await this.tokenBucket.removeTokens(count);
this.tokensThisInterval += count;
return remainingTokens;
}
}
// Remove the requested number of tokens from the token bucket
const remainingTokens = await this.tokenBucket.removeTokens(count);
this.tokensThisInterval += count;
return remainingTokens;
}
/**
* Attempt to remove the requested number of tokens and return immediately.
* If the bucket (and any parent buckets) contains enough tokens and we
* haven't spent too many tokens in this interval already, this will return
* true. Otherwise, false is returned.
* @param {Number} count The number of tokens to remove.
* @param {Boolean} True if the tokens were successfully removed, otherwise
* false.
*/
tryRemoveTokens(count) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize)
return false;
const now = (0, clock_js_1.getMilliseconds)();
// Advance the current interval and reset the current interval token count
// if needed
if (now < this.curIntervalStart || now - this.curIntervalStart >= this.tokenBucket.interval) {
this.curIntervalStart = now;
this.tokensThisInterval = 0;
}
// If we don't have enough tokens left in this interval, return false
if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval)
return false;
// Try to remove the requested number of tokens from the token bucket
const removed = this.tokenBucket.tryRemoveTokens(count);
if (removed) {
this.tokensThisInterval += count;
}
return removed;
}
/**
* Returns the number of tokens remaining in the TokenBucket.
* @returns {Number} The number of tokens remaining.
*/
getTokensRemaining() {
this.tokenBucket.drip();
return this.tokenBucket.content;
}
}
exports.RateLimiter = RateLimiter;
//# sourceMappingURL=RateLimiter.js.map