From 3f864813e7aed98c9d8b2e8b7e75f2c5e1099972 Mon Sep 17 00:00:00 2001 From: JT Date: Wed, 2 Sep 2015 18:32:17 +0100 Subject: [PATCH 1/2] Make a bit of a mess of TokenBucket.lua. I think this works. Also (hopefully) caters for non-monotonic time. Possibly. --- lib/TokenBucket.lua | 55 ++++++++++++++++++++++++++++++--------------- 1 file changed, 37 insertions(+), 18 deletions(-) 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) From 582bca361f4a841677fea3e1677b697daddb6de0 Mon Sep 17 00:00:00 2001 From: Ben White Date: Fri, 4 Sep 2015 16:06:23 +0100 Subject: [PATCH 2/2] Removed 'wihtoutlua' function as it is unnecessary. Removed default poolMax and fillRate props as they do not play well if the Hyacinth is used asynchronously (as is its intention) Changed rateLimit function to accept the fill and max props Updated tests --- README.md | 6 ++-- lib/TokenBucket.js | 70 +++++++--------------------------------- test/TokenBucket.spec.js | 51 +++++++---------------------- 3 files changed, 25 insertions(+), 102 deletions(-) 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/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); })); };