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