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