rate-limiter.js 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. exports.RateLimiter = void 0;
  4. /*!
  5. * Copyright 2020 Google LLC
  6. *
  7. * Licensed under the Apache License, Version 2.0 (the "License");
  8. * you may not use this file except in compliance with the License.
  9. * You may obtain a copy of the License at
  10. *
  11. * http://www.apache.org/licenses/LICENSE-2.0
  12. *
  13. * Unless required by applicable law or agreed to in writing, software
  14. * distributed under the License is distributed on an "AS IS" BASIS,
  15. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16. * See the License for the specific language governing permissions and
  17. * limitations under the License.
  18. */
  19. const assert = require("assert");
  20. const logger_1 = require("./logger");
  21. /**
  22. * A helper that uses the Token Bucket algorithm to rate limit the number of
  23. * operations that can be made in a second.
  24. *
  25. * Before a given request containing a number of operations can proceed,
  26. * RateLimiter determines doing so stays under the provided rate limits. It can
  27. * also determine how much time is required before a request can be made.
  28. *
  29. * RateLimiter can also implement a gradually increasing rate limit. This is
  30. * used to enforce the 500/50/5 rule
  31. * (https://firebase.google.com/docs/firestore/best-practices#ramping_up_traffic).
  32. *
  33. * @private
  34. * @internal
  35. */
  36. class RateLimiter {
  37. /**
  38. * @param initialCapacity Initial maximum number of operations per second.
  39. * @param multiplier Rate by which to increase the capacity.
  40. * @param multiplierMillis How often the capacity should increase in
  41. * milliseconds.
  42. * @param maximumCapacity Maximum number of allowed operations per second.
  43. * The number of tokens added per second will never exceed this number.
  44. * @param startTimeMillis The starting time in epoch milliseconds that the
  45. * rate limit is based on. Used for testing the limiter.
  46. */
  47. constructor(initialCapacity, multiplier, multiplierMillis, maximumCapacity, startTimeMillis = Date.now()) {
  48. this.initialCapacity = initialCapacity;
  49. this.multiplier = multiplier;
  50. this.multiplierMillis = multiplierMillis;
  51. this.maximumCapacity = maximumCapacity;
  52. this.startTimeMillis = startTimeMillis;
  53. this.availableTokens = initialCapacity;
  54. this.lastRefillTimeMillis = startTimeMillis;
  55. this.previousCapacity = initialCapacity;
  56. }
  57. /**
  58. * Tries to make the number of operations. Returns true if the request
  59. * succeeded and false otherwise.
  60. *
  61. * @param requestTimeMillis The time used to calculate the number of available
  62. * tokens. Used for testing the limiter.
  63. * @private
  64. * @internal
  65. */
  66. tryMakeRequest(numOperations, requestTimeMillis = Date.now()) {
  67. this.refillTokens(requestTimeMillis);
  68. if (numOperations <= this.availableTokens) {
  69. this.availableTokens -= numOperations;
  70. return true;
  71. }
  72. return false;
  73. }
  74. /**
  75. * Returns the number of ms needed to make a request with the provided number
  76. * of operations. Returns 0 if the request can be made with the existing
  77. * capacity. Returns -1 if the request is not possible with the current
  78. * capacity.
  79. *
  80. * @param requestTimeMillis The time used to calculate the number of available
  81. * tokens. Used for testing the limiter.
  82. * @private
  83. * @internal
  84. */
  85. getNextRequestDelayMs(numOperations, requestTimeMillis = Date.now()) {
  86. this.refillTokens(requestTimeMillis);
  87. if (numOperations < this.availableTokens) {
  88. return 0;
  89. }
  90. const capacity = this.calculateCapacity(requestTimeMillis);
  91. if (capacity < numOperations) {
  92. return -1;
  93. }
  94. const requiredTokens = numOperations - this.availableTokens;
  95. return Math.ceil((requiredTokens * 1000) / capacity);
  96. }
  97. /**
  98. * Refills the number of available tokens based on how much time has elapsed
  99. * since the last time the tokens were refilled.
  100. *
  101. * @param requestTimeMillis The time used to calculate the number of available
  102. * tokens. Used for testing the limiter.
  103. * @private
  104. * @internal
  105. */
  106. refillTokens(requestTimeMillis) {
  107. if (requestTimeMillis >= this.lastRefillTimeMillis) {
  108. const elapsedTime = requestTimeMillis - this.lastRefillTimeMillis;
  109. const capacity = this.calculateCapacity(requestTimeMillis);
  110. const tokensToAdd = Math.floor((elapsedTime * capacity) / 1000);
  111. if (tokensToAdd > 0) {
  112. this.availableTokens = Math.min(capacity, this.availableTokens + tokensToAdd);
  113. this.lastRefillTimeMillis = requestTimeMillis;
  114. }
  115. }
  116. else {
  117. throw new Error('Request time should not be before the last token refill time.');
  118. }
  119. }
  120. /**
  121. * Calculates the maximum capacity based on the provided date.
  122. *
  123. * @private
  124. * @internal
  125. */
  126. // Visible for testing.
  127. calculateCapacity(requestTimeMillis) {
  128. assert(requestTimeMillis >= this.startTimeMillis, 'startTime cannot be after currentTime');
  129. const millisElapsed = requestTimeMillis - this.startTimeMillis;
  130. const operationsPerSecond = Math.min(Math.floor(Math.pow(this.multiplier, Math.floor(millisElapsed / this.multiplierMillis)) * this.initialCapacity), this.maximumCapacity);
  131. if (operationsPerSecond !== this.previousCapacity) {
  132. (0, logger_1.logger)('RateLimiter.calculateCapacity', null, `New request capacity: ${operationsPerSecond} operations per second.`);
  133. }
  134. this.previousCapacity = operationsPerSecond;
  135. return operationsPerSecond;
  136. }
  137. }
  138. exports.RateLimiter = RateLimiter;
  139. //# sourceMappingURL=rate-limiter.js.map