diff --git a/README.md b/README.md index fcc763b..d19da9a 100644 --- a/README.md +++ b/README.md @@ -137,6 +137,20 @@ rate: '100/h' *Note:* the rate is parsed ahead of time, so this notation doesn't affect performance. +### `incr` + +How much to increment the usage for a given call. This is useful in scenarios where incrementing by +one is not the whole story. For example, for rate limiting computation time. Some operations take +longer than other so counting operations doesn't work. Rather, use `incr` to track and limit the +actual compute time. + +You would typically specify a custom function: + +```js +// rate-limit based on the cost of an operation +key: function(x) { return x.cost; } +``` + ## HTTP middleware This package contains a pre-built middleware, diff --git a/lib/options.js b/lib/options.js index ac547b0..f093159 100644 --- a/lib/options.js +++ b/lib/options.js @@ -22,7 +22,12 @@ var getLimit = function(opts){ return opts.limit; }; -var getWindow = function(opts){ +var getIncr = function (opts) { + if(typeof opts.incr === 'function') return opts.incr; + return function () { return opts.incr || 1; }; +}; + +var getWindow = function (opts) { if(getRate(opts)) return moment.duration(1, getMatches(opts)[2]) / 1000; if(typeof opts.window === 'function') return opts.window(); return opts.window; @@ -38,6 +43,7 @@ var validate = function(opts){ assert.equal(typeof getKey(opts), 'function', 'Invalid key: ' + opts.key); if(opts.rate) assert.ok(getMatches(opts), 'Invalid rate: ' + getRate(opts)); assert.equal(typeof getLimit(opts), 'number', 'Invalid limit: ' + getLimit(opts)); + assert.equal(typeof getIncr(opts), 'function', 'Invalid incr: ' + opts.incr); assert.equal(typeof getWindow(opts), 'number', 'Invalid window: ' + getWindow(opts)); assert.notEqual(getLimit(opts), 0, 'Invalid rate limit: ' + getRate(opts)); assert.notEqual(getWindow(opts), 0, 'Invalid rate window: ' + getRate(opts)); @@ -50,6 +56,7 @@ canonical = function(opts) { key: getKey(opts), rate: getRate.bind(null, opts), limit: getLimit.bind(null, opts), + incr: getIncr(opts), window: getWindow.bind(null, opts) }; }; diff --git a/lib/rate-limiter.js b/lib/rate-limiter.js index e022841..c572059 100644 --- a/lib/rate-limiter.js +++ b/lib/rate-limiter.js @@ -4,12 +4,13 @@ module.exports = function(opts) { opts = options.canonical(opts); return function(request, callback) { var key = opts.key(request); + var incr = opts.incr(request); var tempKey = 'ratelimittemp:' + key; var realKey = 'ratelimit:' + key; opts.redis.multi() .setex(tempKey, opts.window(), 0) .renamenx(tempKey, realKey) - .incr(realKey) + .incrby(realKey, incr) .ttl(realKey) .exec(function(err, results) { if(err) {