'use strict';
const TokenBucket = require('./TokenBucket');
const _ = require('lodash');
/**
* RateLimiter Class
*
* Handles creation and destruction of TokenBuckets by key.
*/
class RateLimiter {
/**
* Creates a RateLimiter object.
*
* @param {object} config configuration values
* @param {integer} config.limit The max amount of tokens per incrementInterval.
* @param {integer} config.incrementInterval The time in ms before tokens will increment.
* @param {integer} config.increment The amount of tokens that will increment each incrementInterval.
*/
constructor(config) {
let limit = null;
let incrementInterval = null;
let increment = null;
if (config != null) {
limit = config.limit;
incrementInterval = config.incrementInterval;
increment = config.increment;
}
this.config = { limit, incrementInterval, increment };
this.tokenBuckets = {};
this._interval = setInterval(() => {
this._cleanup();
}, incrementInterval);
}
/**
* Cleans up unused TokenBuckets
*/
_cleanup() {
_.each(this.tokenBuckets, (value, key) => {
const tokenBucket = this.tokenBuckets[key];
if (tokenBucket.tokens == tokenBucket.limit) {
tokenBucket.destroy();
delete this.tokenBuckets[key];
}
});
}
/**
* Cleans up all TokenBuckets and nulls the array.
*/
destroy() {
_.each(this.tokenBuckets, (value, key) => {
this.tokenBuckets[key].destroy();
delete this.tokenBuckets[key];
});
}
/**
* Helper to throw an error if the key is invalid.
*
* @param {key} key The identifier for the request
* @throws {Error} throw error when key is invalid
* @private
*/
_errorCheckForKey(key) {
if (key == null) {
throw new Error('key is required');
}
else if(typeof(key) != 'string') {
throw new Error('key is not a string');
}
}
/**
* Helper to create a new TokenBucket if needed.
*
* @param {string} key The identifier for the request
* @private
*/
_getTokenBucket(key) {
this._errorCheckForKey(key);
if (this.tokenBuckets[key] == null) {
const limit = this.config.limit;
const increment = this.config.increment;
const incrementInterval = this.config.incrementInterval;
this.tokenBuckets[key] = new TokenBucket({ limit, increment, incrementInterval });
}
return this.tokenBuckets[key];
}
/**
* Appropriately decrements the remaining token availability
* when called and returns the remainder.
*
* @param {string} key The identifier for the request
* @param {integer} [amount] The number to decrement by
* @return {integer} Remaining request count
*/
decrementTokens(key, amount) {
const tokenBucket = this._getTokenBucket(key);
return tokenBucket.decrementTokens(amount);
}
/**
* Returns the remaining request count for the given key.
* If the key doesn't already exist, it will create it.
*
* @param {string} key The identifier for the request
* @return {integer} Remaining request count
*/
getTokensRemaining(key) {
const tokenBucket = this._getTokenBucket(key);
return tokenBucket.getTokensRemaining();
}
}
module.exports = RateLimiter;
|