diff --git a/.gitignore b/.gitignore index 10d9b1f..e776630 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /node_modules /dump.rdb .DS_Store +.idea \ No newline at end of file diff --git a/lib/options.js b/lib/options.js index ac547b0..977dfef 100644 --- a/lib/options.js +++ b/lib/options.js @@ -50,7 +50,9 @@ canonical = function(opts) { key: getKey(opts), rate: getRate.bind(null, opts), limit: getLimit.bind(null, opts), - window: getWindow.bind(null, opts) + window: getWindow.bind(null, opts), + deleteImmediatelyIfRaceCondition: opts.deleteImmediatelyIfRaceCondition, + onPossibleRaceCondition: opts.onPossibleRaceCondition }; }; diff --git a/lib/rate-limiter.js b/lib/rate-limiter.js index e022841..7e977b4 100644 --- a/lib/rate-limiter.js +++ b/lib/rate-limiter.js @@ -15,10 +15,21 @@ module.exports = function(opts) { if(err) { callback(err); } else { - if (results[3] == -1) { // automatically recover from possible race condition - opts.redis.expire(realKey,opts.window()); + // if multi() used, ioredis returns an array of result set [err | null, value], we need to check the second parameter for such cases + // see also: https://github.com/luin/ioredis/wiki/Migrating-from-node_redis + const hasTimeToLive = Array.isArray(results[3]) ? results[3][1] : results[3]; + if (hasTimeToLive === -1) { // automatically recover from possible race condition + if (opts.deleteImmediatelyIfRaceCondition) { + opts.redis.del(realKey); + } else { + opts.redis.expire(realKey, opts.window()); + } + + if (typeof opts.onPossibleRaceCondition === 'function') { + opts.onPossibleRaceCondition(realKey); + } } - var current = results[2]; + const current = Array.isArray(results[2]) ? results[2][1] : results[2]; callback(null, { key: key, current: current, diff --git a/package.json b/package.json index 9b62d3e..9e10c21 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,7 @@ "devDependencies": { "async": "^2.6.0", "express": "^4.16.2", + "ioredis": "^4.2.0", "lodash": "^4.17.4", "mocha": "^4.0.1", "redis": "^2.8.0", diff --git a/test/rate-limiter.ioredis.spec.js b/test/rate-limiter.ioredis.spec.js new file mode 100644 index 0000000..08c3b26 --- /dev/null +++ b/test/rate-limiter.ioredis.spec.js @@ -0,0 +1,123 @@ +const _ = require('lodash'); +const async = require('async'); +const redis = require('ioredis'); +const reset = require('./reset'); +const rateLimiter = require('../lib/rate-limiter'); + +describe('IORedis Client test Rate-limiter', function () { + + this.slow(5000); + this.timeout(5000); + + let client = null; + + before(function (done) { + client = new redis(6379, 'localhost', {enable_offline_queue: false}); + client.on('ready', done); + }); + + beforeEach(function (done) { + reset.allkeys(client, done); + }); + + after(function () { + client.quit(); + }); + + it('calls back with the rate data', function (done) { + var limiter = createLimiter('10/second'); + var reqs = request(limiter, 5, {id: 'a'}); + async.parallel(reqs, function (err, rates) { + _.map(rates, 'current').should.eql([1, 2, 3, 4, 5]); + _.each(rates, function (r) { + r.key.should.eql('a'); + r.limit.should.eql(10); + r.window.should.eql(1); + r.over.should.eql(false); + }); + done(); + }); + }); + + it('sets the over flag when above the limit', function (done) { + var limiter = createLimiter('10/second'); + var reqs = request(limiter, 15, {id: 'a'}); + async.parallel(reqs, function (err, rates) { + _.each(rates, function (r, index) { + rates[index].over.should.eql(index >= 10); + }); + done(); + }); + }); + + it('uses one bucket per key', function (done) { + var limiter = createLimiter('10/second'); + var reqs = _.flatten([ + request(limiter, 10, {id: 'a'}), + request(limiter, 12, {id: 'b'}), + request(limiter, 10, {id: 'c'}) + ]); + async.parallel(reqs, function (err, rates) { + _.filter(rates, {over: true}).should.have.length(2); + done(); + }); + }); + + it('can handle a lot of requests', function (done) { + var limiter = createLimiter('1000/second'); + var reqs = request(limiter, 1200, {id: 'a'}); + async.parallel(reqs, function (err, rates) { + rates[999].should.have.property('over', false); + rates[1000].should.have.property('over', true); + done(); + }); + }); + + it('resets after the window', function (done) { + var limiter = createLimiter('10/second'); + async.series([ + requestParallel(limiter, 15, {id: 'a'}), + wait(1100), + requestParallel(limiter, 15, {id: 'a'}) + ], function (err, data) { + _.each(data[0], function (rate, index) { + rate.should.have.property('over', index > 9); + }); + _.each(data[2], function (rate, index) { + rate.should.have.property('over', index > 9); + }); + done(); + }); + }); + + function createLimiter(rate) { + return rateLimiter({ + redis: client, + key: function (x) { + return x.id + }, + rate: rate + }); + } + + function request(limiter, count, data) { + return _.times(count, function () { + return function (next) { + limiter(data, next); + }; + }); + } + + function requestParallel(limiter, count, data) { + return function (next) { + async.parallel(request(limiter, count, data), next); + }; + } + + function wait(millis) { + return function (next) { + setTimeout(next, 1100); + }; + } + +});