work.suroh.tk/node_modules/limiter/lib/rateLimiter.js

138 lines
4.8 KiB
JavaScript
Raw Normal View History

2019-12-02 12:22:45 +00:00
var TokenBucket = require('./tokenBucket');
var getMilliseconds = require('./clock');
/**
* 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.
* @author John Hurliman <jhurliman@jhurliman.org>
*
* @param {Number} tokensPerInterval Maximum number of tokens that can be
* removed at any given moment and over the course of one interval.
* @param {String|Number} interval The interval length in milliseconds, or as
* one of the following strings: 'second', 'minute', 'hour', day'.
* @param {Boolean} fireImmediately Optional. Whether or not the callback
* will fire immediately when rate limiting is in effect (default is false).
*/
var RateLimiter = function(tokensPerInterval, interval, fireImmediately) {
this.tokenBucket = new TokenBucket(tokensPerInterval, tokensPerInterval,
interval, null);
// Fill the token bucket to start
this.tokenBucket.content = tokensPerInterval;
this.curIntervalStart = getMilliseconds();
this.tokensThisInterval = 0;
this.fireImmediately = fireImmediately;
};
RateLimiter.prototype = {
tokenBucket: null,
curIntervalStart: 0,
tokensThisInterval: 0,
fireImmediately: false,
/**
* Remove the requested number of tokens and fire the given callback. 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 and callback will happen when enough tokens become available.
* @param {Number} count The number of tokens to remove.
* @param {Function} callback(err, remainingTokens)
* @returns {Boolean} True if the callback was fired immediately, otherwise
* false.
*/
removeTokens: function(count, callback) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize) {
process.nextTick(callback.bind(null, 'Requested tokens ' + count +
' exceeds maximum tokens per interval ' + this.tokenBucket.bucketSize,
null));
return false;
}
var self = this;
var now = 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) {
process.nextTick(callback.bind(null, null, -1));
} else {
var waitInterval = Math.ceil(
this.curIntervalStart + this.tokenBucket.interval - now);
setTimeout(function() {
self.tokenBucket.removeTokens(count, afterTokensRemoved);
}, waitInterval);
}
return false;
}
// Remove the requested number of tokens from the token bucket
return this.tokenBucket.removeTokens(count, afterTokensRemoved);
function afterTokensRemoved(err, tokensRemaining) {
if (err) return callback(err, null);
self.tokensThisInterval += count;
callback(null, tokensRemaining);
}
},
/**
* 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: function(count) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize)
return false;
var now = 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
var 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: function () {
this.tokenBucket.drip();
return this.tokenBucket.content;
}
};
module.exports = RateLimiter;