Subtheme of barrio
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.
 
 
 
 

167 lines
5.6 KiB

/**
* A hierarchical token bucket for rate limiting. See
* http://en.wikipedia.org/wiki/Token_bucket for more information.
* @author John Hurliman <jhurliman@cull.tv>
*
* @param {Number} bucketSize Maximum number of tokens to hold in the bucket.
* Also known as the burst rate.
* @param {Number} tokensPerInterval Number of tokens to drip into the bucket
* 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 {TokenBucket} parentBucket Optional. A token bucket that will act as
* the parent of this bucket.
*/
var TokenBucket = function(bucketSize, tokensPerInterval, interval, parentBucket) {
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('Invaid interval ' + interval);
}
} else {
this.interval = interval;
}
this.parentBucket = parentBucket;
this.content = 0;
this.lastDrip = +new Date();
};
TokenBucket.prototype = {
bucketSize: 1,
tokensPerInterval: 1,
interval: 1000,
parentBucket: null,
content: 0,
lastDrip: 0,
/**
* Remove the requested number of tokens and fire the given callback. If the
* bucket (and any parent buckets) contains enough tokens 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) {
var self = this;
// Is this an infinite size bucket?
if (!this.bucketSize) {
process.nextTick(callback.bind(null, null, count, Number.POSITIVE_INFINITY));
return true;
}
// Make sure the bucket can hold the requested number of tokens
if (count > this.bucketSize) {
process.nextTick(callback.bind(null, 'Requested tokens ' + count + ' exceeds bucket size ' +
this.bucketSize, null));
return false;
}
// Drip new tokens into this bucket
this.drip();
// If we don't have enough tokens in this bucket, come back later
if (count > this.content)
return comeBackLater();
if (this.parentBucket) {
// Remove the requested from the parent bucket first
return this.parentBucket.removeTokens(count, function(err, remainingTokens) {
if (err) return callback(err, null);
// Check that we still have enough tokens in this bucket
if (count > self.content)
return comeBackLater();
// Tokens were removed from the parent bucket, now remove them from
// this bucket and fire the callback. Note that we look at the current
// bucket and parent bucket's remaining tokens and return the smaller
// of the two values
self.content -= count;
callback(null, Math.min(remainingTokens, self.content));
});
} else {
// Remove the requested tokens from this bucket and fire the callback
this.content -= count;
process.nextTick(callback.bind(null, null, this.content));
return true;
}
function comeBackLater() {
// How long do we need to wait to make up the difference in tokens?
var waitInterval = Math.ceil(
(count - self.content) * (self.interval / self.tokensPerInterval));
setTimeout(function() { self.removeTokens(count, callback); }, waitInterval);
return false;
}
},
/**
* 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: function(count) {
// Is this an infinite size bucket?
if (!this.bucketSize)
return true;
// Make sure the bucket can hold the requested number of tokens
if (count > this.bucketSize)
return false;
// Drip new tokens into this bucket
this.drip();
// If we don't have enough tokens in this bucket, return false
if (count > this.content)
return false;
// Try to remove the requested tokens from the parent bucket
if (this.parentBucket && !this.parentBucket.tryRemoveTokens(count))
return false;
// Remove the requested tokens from this bucket and return
this.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: function() {
if (!this.tokensPerInterval) {
this.content = this.bucketSize;
return;
}
var now = +new Date();
var deltaMS = Math.max(now - this.lastDrip, 0);
this.lastDrip = now;
var dripAmount = deltaMS * (this.tokensPerInterval / this.interval);
this.content = Math.min(this.content + dripAmount, this.bucketSize);
}
};
module.exports = TokenBucket;