Skip to content

Commit

Permalink
Merge pull request #3 from gosquared/pruning-and-improving
Browse files Browse the repository at this point in the history 
Pruning and improving
  • Loading branch information
@benjackwhite
benjackwhite committed Sep 4, 2015
2 parents 7952636 + 582bca3 commit 969b9e1
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 120 deletions.
6 changes: 2 additions & 4 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
70 changes: 12 additions & 58 deletions lib/TokenBucket.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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.<number, Error>} - 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);
Expand All @@ -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.<number,Error} The promise object for this operation
*/
TokenBucket.prototype.rateLimitWithoutLua = function(key, cost, cb) {
var self = this;
cb = cb || function() {};

return new Promise(function(resolve, reject) {

var now = Date.now();

// Get the last timestamp to compare
self.redis.getset(key + 'timestamp', now, function(err, res){

if(!res) {
// There is no timestamp to compare set the poolmax and
// respond with that minus cost
self.redis.set(key + 'pool', self.poolMax - cost, function(err, res){
});

resolve(self.poolMax - cost);
return;
}

// Calculate the difference (amount owed minus cost)
var owed = (now - res) / self.fillRate;

// Get the current pool amount to compare
self.redis.get(key + 'pool', function(err, res){
// Here we want to add on the owed tokens up to the poolMax
res = (res === null) ? self.poolMax : parseInt(res, 10);

var beforeCost = (res + owed < self.poolMax) ? res + owed : self.poolMax;
var afterCost = beforeCost - cost;
var newAmount = (afterCost >= 0) ? afterCost : beforeCost;
var limited = (afterCost >= 0) ? 1 : -1;

self.redis.set(key + 'pool', newAmount, function(){
return resolve(newAmount * limited);
});
});
});
});
};
55 changes: 37 additions & 18 deletions lib/TokenBucket.lua
View file
Original file line number Diff line number Diff line change
@@ -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)
51 changes: 11 additions & 40 deletions test/TokenBucket.spec.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ var timers = 0;
var count = 0;
var client;


describe('TokenBucket', function(){

before(function(done){
Expand All @@ -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;
});
Expand All @@ -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);
});
});
Expand All @@ -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);
});
});
Expand All @@ -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);
});
});
Expand All @@ -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);
});
});
Expand All @@ -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);
}));
};
Expand Down

0 comments on commit 969b9e1

Please sign in to comment.