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
183 lines • 7.14 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.TokenBucket = void 0;
exports.validateTokens = validateTokens;
const RequestQueue_js_1 = require("./RequestQueue.js");
const clock_js_1 = require("./clock.js");
/**
* A hierarchical token bucket for rate limiting. See
* http://en.wikipedia.org/wiki/Token_bucket for more information.
*
* @param options
* @param options.bucketSize Maximum number of tokens to hold in the bucket.
* Also known as the burst rate.
* @param options.tokensPerInterval Number of tokens to drip into the bucket
* 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.parentBucket Optional. A token bucket that will act as
* the parent of this bucket.
*/
class TokenBucket {
requests = new RequestQueue_js_1.RequestQueue();
bucketSize;
tokensPerInterval;
interval;
parentBucket;
content;
lastDrip;
constructor({ bucketSize, tokensPerInterval, interval, parentBucket }) {
validateTokens(bucketSize, "bucketSize");
validateTokens(tokensPerInterval, "tokensPerInterval");
this.bucketSize = bucketSize;
this.tokensPerInterval = tokensPerInterval;
if (typeof interval === "string") {
switch (interval) {
case "sec":
case "second":
this.interval = 1000;
break;
case "min":
case "minute":
this.interval = 1000 * 60;
break;
case "hr":
case "hour":
this.interval = 1000 * 60 * 60;
break;
case "day":
this.interval = 1000 * 60 * 60 * 24;
break;
default:
throw new Error("Invalid interval " + interval);
}
}
else {
this.interval = interval;
}
if (!Number.isFinite(this.interval) || this.interval <= 0) {
throw new RangeError("interval must be a finite positive number of milliseconds");
}
this.parentBucket = parentBucket;
this.content = 0;
this.lastDrip = (0, clock_js_1.getMilliseconds)();
}
/**
* Remove the requested number of tokens. If the bucket (and any parent
* buckets) contains enough tokens 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) {
validateTokens(count, "count");
this.getWaitTime(count);
return this.requests.run(async () => {
while (true) {
if (this.tryRemoveTokens(count))
return this.remainingTokens();
await (0, clock_js_1.wait)(Math.max(1, this.getWaitTime(count)));
}
});
}
/** Return the wait required by this bucket and all of its ancestors. */
getWaitTime(count) {
validateTokens(count, "count");
let delay = 0;
for (const bucket of hierarchy(this)) {
if (count > bucket.bucketSize) {
throw new RangeError(`Requested tokens ${count} exceeds bucket size ${bucket.bucketSize}`);
}
bucket.drip();
if (count > bucket.content) {
delay = Math.max(delay, Math.ceil((count - bucket.content) / (bucket.tokensPerInterval / bucket.interval)));
}
}
return delay;
}
remainingTokens() {
if (!this.parentBucket)
return this.bucketSize === 0 ? Number.POSITIVE_INFINITY : this.content;
let remaining = Number.POSITIVE_INFINITY;
for (const bucket of hierarchy(this)) {
remaining = Math.min(remaining, bucket.content);
}
return remaining;
}
/**
* Attempt to remove the requested number of tokens and return immediately.
* If the bucket (and any parent buckets) contains enough tokens 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) {
validateTokens(count, "count");
// Check the whole hierarchy before charging any bucket. No await may occur
// between this check and the debit: competing callers must see the debit.
if (this.bucketSize === 0)
return true;
if (!this.parentBucket) {
if (count > this.bucketSize)
return false;
this.drip();
if (count > this.content)
return false;
this.content -= count;
return true;
}
const buckets = hierarchy(this);
for (const bucket of buckets) {
if (count > bucket.bucketSize)
return false;
bucket.drip();
if (count > bucket.content)
return false;
}
for (const bucket of buckets)
bucket.content -= count;
return true;
}
/**
* Add any new tokens to the bucket since the last drip.
* @returns {Boolean} True if new tokens were added, otherwise false.
*/
drip() {
if (this.tokensPerInterval === 0) {
const prevContent = this.content;
this.content = this.bucketSize;
return this.content > prevContent;
}
const now = (0, clock_js_1.getMilliseconds)();
const deltaMS = Math.max(now - this.lastDrip, 0);
this.lastDrip = now;
// An extremely short interval can overflow the rate even though both
// inputs are finite. In particular, 0 * Infinity would poison the bucket.
const rate = this.tokensPerInterval / this.interval;
const dripAmount = Number.isFinite(rate)
? deltaMS * rate
: (deltaMS * this.tokensPerInterval) / this.interval;
const prevContent = this.content;
this.content = Math.min(this.content + dripAmount, this.bucketSize);
return Math.floor(this.content) > Math.floor(prevContent);
}
}
exports.TokenBucket = TokenBucket;
function validateTokens(value, name) {
if (!Number.isFinite(value) || value < 0 || value > Number.MAX_SAFE_INTEGER) {
throw new RangeError(`${name} must be a finite non-negative number no greater than Number.MAX_SAFE_INTEGER`);
}
}
function hierarchy(start) {
const buckets = [];
const visited = new Set();
for (let bucket = start; bucket && bucket.bucketSize !== 0; bucket = bucket.parentBucket) {
if (visited.has(bucket))
throw new Error("Circular parentBucket hierarchy");
visited.add(bucket);
buckets.push(bucket);
}
return buckets;
}
//# sourceMappingURL=TokenBucket.js.map