diff --git a/README.md b/README.md index 6d83251..aef777d 100644 --- a/README.md +++ b/README.md @@ -16,12 +16,10 @@ npm install hyacinth var TokenBucket = require('hyacinth'); var rateLimiter = new TokenBucket({ - redis: redisClient, - poolMax: 250, - fillRate: 240, + redis: redisClient }); -rateLimiter.rateLimit(testKey, 10).then(function(tokensRemaining){ +rateLimiter.rateLimit(testKey, 10, 250, 240).then(function(tokensRemaining){ // Negative number indicates the tokens remaining but limited // as the cost was higher than those remaining diff --git a/lib/TokenBucket.js b/lib/TokenBucket.js index 5a2b644..77f6a34 100644 --- a/lib/TokenBucket.js +++ b/lib/TokenBucket.js @@ -9,9 +9,6 @@ var Scripty = require('node-redis-scripty'); * Constructor for TokenBucket * @param {Object} config - The configuration object for the TokenBucket * @param {Object} config.redis - The redis client to be used as the store - * @param {number} [config.poolMax=250] - The maximum size of the token pool - * @param {number} [config.fillRate=240] - The rate in milliseconds that the - * pool will be refilled * @param {string} [config.identifier=api-token-bucket-] - The identifier to * be prepended to all redis keys * @return {TokenBucket} A new TokenBucket instance @@ -21,8 +18,6 @@ var TokenBucket = module.exports = function(config) { if (!config) config = {}; this.redis = config.redis || Redis.createClient(); - this.poolMax = config.poolMax || 250; - this.fillRate = config.fillRate || 240; //milliseconds between fill this.scripty = new Scripty(this.redis); return this; @@ -55,21 +50,29 @@ TokenBucket.prototype.clearRateLimitWithKey = function(key, cb) { * Calculates the rate limit with a call to redis for a given key * @param {string} key The key to be rate limited * @param {number} cost The cost of the operation - * @param {Function} cb The callback function + * @param {number} poolMax The max value for the pool + * @param {number} fillRate The fill rate for the pool + * @param {Function} [cb] The callback function * @return {Promise.} - Resolves to the number of tokens left with * a negative number indicating that the req was limited */ -TokenBucket.prototype.rateLimit = function(key, cost, cb) { +TokenBucket.prototype.rateLimit = function(key, cost, poolMax, fillRate, cb) { var self = this; + cb = cb || function() {}; + if (!key || !cost || !poolMax || !fillRate) { + cb(new Error('Missing arguments')); + return Promise.reject(new Error('Missing arguments')); + } + return new Promise(function(resolve, reject) { var now = Date.now(); self.scripty.loadScriptFile('blank', __dirname + '/TokenBucket.lua', function(err, script) { - if(err) throw new Error(err); - script.run(6, key, now, cost, self.poolMax, self.fillRate, Math.ceil(self.fillRate * self.poolMax / 1000), function(err, res) { + if (err) throw new Error(err); + script.run(6, key, now, cost, poolMax, fillRate, Math.ceil(fillRate * poolMax / 1000), function(err, res) { if (err) { cb(err, null); reject(err); @@ -82,52 +85,3 @@ TokenBucket.prototype.rateLimit = function(key, cost, cb) { }); }); }; - -/** - * Calculates the rate limit with a calls to Redis but in a non-atomic way without Lua. Should be used as a reference rather than an actual implementation - * @param {string} key The key to be rate limited - * @param {number} cost The cost of the operation - * @param {Function} cb The callback function - * @return {Promise.= 0) ? afterCost : beforeCost; - var limited = (afterCost >= 0) ? 1 : -1; - - self.redis.set(key + 'pool', newAmount, function(){ - return resolve(newAmount * limited); - }); - }); - }); - }); -}; diff --git a/lib/TokenBucket.lua b/lib/TokenBucket.lua index 6576f36..c429b54 100644 --- a/lib/TokenBucket.lua +++ b/lib/TokenBucket.lua @@ -1,28 +1,47 @@ -local t = redis.call('getset',KEYS[1]..'timestamp', KEYS[2]) -redis.call('expire', KEYS[1]..'timestamp', KEYS[6]) -if t == false then - redis.call('set', KEYS[1]..'pool', KEYS[4] - KEYS[3]) - redis.call('expire',KEYS[1]..'pool', KEYS[6]) - return tostring(KEYS[4] - KEYS[3]) -end +local key = KEYS[1] +local now = KEYS[2] +local cost = KEYS[3] +local poolMax = tonumber(KEYS[4]) +local fillRate = KEYS[5] +local expiry = KEYS[6] -local owed = (KEYS[2] - t) / KEYS[5] -local r = redis.call('get', KEYS[1]..'pool') +local timestampKey = key..'timestamp' +local poolKey = key..'pool' -if r == false then - r = KEYS[4] +local before = redis.call('get', timestampKey) + +if before == false then + redis.call('set', timestampKey, now, 'ex', expiry) + + local ret = poolMax - cost + redis.call('set', poolKey, ret, 'ex', expiry) + return tostring(ret) end -if r + owed < tonumber(KEYS[4]) then - r = r + owed + +local timediff = now - before + +if timediff > 0 then + redis.call('set', timestampKey, now, 'ex', expiry) else - r = KEYS[4] + timediff = 0 +end + +local owed = timediff / fillRate +local r = redis.call('get', poolKey) + +if r == false then + r = poolMax end + +r = math.min(r + owed, poolMax) + local limit = 1 -if r - KEYS[3] >= 0 then - r = r - KEYS[3] +if r - cost >= 0 then + r = r - cost else limit = -1 end -redis.call('set', KEYS[1]..'pool', r) -redis.call('expire',KEYS[1]..'pool', KEYS[6]) + +redis.call('set', poolKey, r, 'ex', expiry) + return tostring(r * limit) diff --git a/test/TokenBucket.spec.js b/test/TokenBucket.spec.js index 669fc85..cba3b99 100644 --- a/test/TokenBucket.spec.js +++ b/test/TokenBucket.spec.js @@ -9,6 +9,7 @@ var timers = 0; var count = 0; var client; + describe('TokenBucket', function(){ before(function(done){ @@ -19,21 +20,16 @@ describe('TokenBucket', function(){ }); client.on('ready', function(){ + rateLimiter = new TokenBucket({redis:client}); done(); }); }); describe('rateLimitReset', function(){ - it('should return the max pool size when resetting', function(){ + it('should return true when resetting', function(){ var testKey = 'API:limits:testing:0:'; - rateLimiter = new TokenBucket({ - redis:client, - poolMax: 250, - fillRate: 40 - }); - return rateLimiter.clearRateLimitWithKey(testKey).then(function(data){ expect(data).to.be.true; }); @@ -51,9 +47,8 @@ describe('TokenBucket', function(){ it('should return the the pool max minus the cost after being reset', function(){ var testKey = 'API:limits:testing:1:'; - rateLimiter = new TokenBucket({redis:client}); - return rateLimiter.rateLimit(testKey, 10).then(function(data){ + return rateLimiter.rateLimit(testKey, 10, 250, 240).then(function(data){ expect(data).to.equal(240); }); }); @@ -66,7 +61,7 @@ describe('TokenBucket', function(){ this.timeout(4000); return testRateLimit(250, 240, 250, 2000, 1).then(function(data){ - var passed = data.filter(function(item){return item >= 0}).length; + var passed = data.filter(function(item){return item >= 0;}).length; expect(passed).to.equal(250); }); }); @@ -75,7 +70,7 @@ describe('TokenBucket', function(){ this.timeout(4000); return testRateLimit(250, 240, 250, 2000, 1.5).then(function(data){ - var passed = data.filter(function(item){return item >= 0}).length; + var passed = data.filter(function(item){return item >= 0;}).length; expect(passed).to.equal(172); }); }); @@ -84,7 +79,7 @@ describe('TokenBucket', function(){ this.timeout(4000); return testRateLimit(250, 240, 500, 2000, 1).then(function(data){ - var passed = data.filter(function(item){return item >= 0}).length; + var passed = data.filter(function(item){return item >= 0;}).length; expect(passed).to.equal(258); }); }); @@ -93,50 +88,26 @@ describe('TokenBucket', function(){ this.timeout(4000); return testRateLimit(250, 240, 500, 1000, 1).then(function(data){ - var passed = data.filter(function(item){return item >= 0}).length; + var passed = data.filter(function(item){return item >= 0;}).length; expect(passed).to.equal(254); }); }); }); - - describe('rateLimitWithoutLua', function(){ - - beforeEach(function(done){ - client.eval('return redis.call("del", unpack(redis.call("keys", KEYS[1])))', 1, 'API:limits:testing:*', function(err, res){ - done(); - }); - }); - - it('should allow 25 hits out of 500 over 2 seconds at a cost of 10', function(){ - this.timeout(4000); - - return testRateLimit(250, 240, 500, 2000, 10, true).then(function(data){ - var passed = data.filter(function(item){return item >= 0}).length; - expect(passed).to.be.within(24,26); - }); - }); - }); }); -function testRateLimit(poolMax, fillRate, hits, time, cost, withoutLua) { +function testRateLimit(poolMax, fillRate, hits, time, cost) { //Expected pass amount is poolMax + time(ms) / fillRate / cost var promises = []; var key = 'API:limits:testing:'; - rateLimiter = new TokenBucket({ - redis:client, - poolMax: poolMax, - fillRate: fillRate - }); - for(var i =0; i < hits; i += 1){ + for (var i =0; i < hits; i += 1){ promises.push(new Promise(function(resolve, reject) { setTimeout(function(){ - if(withoutLua) rateLimiter.rateLimitWithoutLua(key, cost).then(resolve).catch(reject); - else rateLimiter.rateLimit(key, cost).then(resolve).catch(reject); + rateLimiter.rateLimit(key, cost, poolMax, fillRate).then(resolve).catch(reject); }, (time / hits) * i); })); };