rateLimiter.js 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. var TokenBucket = require('./tokenBucket');
  2. var getMilliseconds = require('./clock');
  3. /**
  4. * A generic rate limiter. Underneath the hood, this uses a token bucket plus
  5. * an additional check to limit how many tokens we can remove each interval.
  6. * @author John Hurliman <jhurliman@jhurliman.org>
  7. *
  8. * @param {Number} tokensPerInterval Maximum number of tokens that can be
  9. * removed at any given moment and over the course of one interval.
  10. * @param {String|Number} interval The interval length in milliseconds, or as
  11. * one of the following strings: 'second', 'minute', 'hour', day'.
  12. * @param {Boolean} fireImmediately Optional. Whether or not the callback
  13. * will fire immediately when rate limiting is in effect (default is false).
  14. */
  15. var RateLimiter = function(tokensPerInterval, interval, fireImmediately) {
  16. this.tokenBucket = new TokenBucket(tokensPerInterval, tokensPerInterval,
  17. interval, null);
  18. // Fill the token bucket to start
  19. this.tokenBucket.content = tokensPerInterval;
  20. this.curIntervalStart = getMilliseconds();
  21. this.tokensThisInterval = 0;
  22. this.fireImmediately = fireImmediately;
  23. };
  24. RateLimiter.prototype = {
  25. tokenBucket: null,
  26. curIntervalStart: 0,
  27. tokensThisInterval: 0,
  28. fireImmediately: false,
  29. /**
  30. * Remove the requested number of tokens and fire the given callback. If the
  31. * rate limiter contains enough tokens and we haven't spent too many tokens
  32. * in this interval already, this will happen immediately. Otherwise, the
  33. * removal and callback will happen when enough tokens become available.
  34. * @param {Number} count The number of tokens to remove.
  35. * @param {Function} callback(err, remainingTokens)
  36. * @returns {Boolean} True if the callback was fired immediately, otherwise
  37. * false.
  38. */
  39. removeTokens: function(count, callback) {
  40. // Make sure the request isn't for more than we can handle
  41. if (count > this.tokenBucket.bucketSize) {
  42. process.nextTick(callback.bind(null, 'Requested tokens ' + count +
  43. ' exceeds maximum tokens per interval ' + this.tokenBucket.bucketSize,
  44. null));
  45. return false;
  46. }
  47. var self = this;
  48. var now = getMilliseconds();
  49. // Advance the current interval and reset the current interval token count
  50. // if needed
  51. if (now < this.curIntervalStart
  52. || now - this.curIntervalStart >= this.tokenBucket.interval) {
  53. this.curIntervalStart = now;
  54. this.tokensThisInterval = 0;
  55. }
  56. // If we don't have enough tokens left in this interval, wait until the
  57. // next interval
  58. if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval) {
  59. if (this.fireImmediately) {
  60. process.nextTick(callback.bind(null, null, -1));
  61. } else {
  62. var waitInterval = Math.ceil(
  63. this.curIntervalStart + this.tokenBucket.interval - now);
  64. setTimeout(function() {
  65. self.tokenBucket.removeTokens(count, afterTokensRemoved);
  66. }, waitInterval);
  67. }
  68. return false;
  69. }
  70. // Remove the requested number of tokens from the token bucket
  71. return this.tokenBucket.removeTokens(count, afterTokensRemoved);
  72. function afterTokensRemoved(err, tokensRemaining) {
  73. if (err) return callback(err, null);
  74. self.tokensThisInterval += count;
  75. callback(null, tokensRemaining);
  76. }
  77. },
  78. /**
  79. * Attempt to remove the requested number of tokens and return immediately.
  80. * If the bucket (and any parent buckets) contains enough tokens and we
  81. * haven't spent too many tokens in this interval already, this will return
  82. * true. Otherwise, false is returned.
  83. * @param {Number} count The number of tokens to remove.
  84. * @param {Boolean} True if the tokens were successfully removed, otherwise
  85. * false.
  86. */
  87. tryRemoveTokens: function(count) {
  88. // Make sure the request isn't for more than we can handle
  89. if (count > this.tokenBucket.bucketSize)
  90. return false;
  91. var now = getMilliseconds();
  92. // Advance the current interval and reset the current interval token count
  93. // if needed
  94. if (now < this.curIntervalStart
  95. || now - this.curIntervalStart >= this.tokenBucket.interval) {
  96. this.curIntervalStart = now;
  97. this.tokensThisInterval = 0;
  98. }
  99. // If we don't have enough tokens left in this interval, return false
  100. if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval)
  101. return false;
  102. // Try to remove the requested number of tokens from the token bucket
  103. var removed = this.tokenBucket.tryRemoveTokens(count);
  104. if (removed) {
  105. this.tokensThisInterval += count;
  106. }
  107. return removed;
  108. },
  109. /**
  110. * Returns the number of tokens remaining in the TokenBucket.
  111. * @returns {Number} The number of tokens remaining.
  112. */
  113. getTokensRemaining: function () {
  114. this.tokenBucket.drip();
  115. return this.tokenBucket.content;
  116. }
  117. };
  118. module.exports = RateLimiter;