You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
137 lines
4.9 KiB
137 lines
4.9 KiB
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;
|
|
|