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
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;
|
|
|