-
Notifications
You must be signed in to change notification settings - Fork 0
/
index.js
100 lines (83 loc) · 2.51 KB
/
index.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
var assert = require("assert");
var fs = require("fs");
var lua = fs.readFileSync("./src/token-bucket.lua");
function isInt(x) {
return x === (x | 0);
}
function RedisRateLimiter(options) {
var redis = options.redis,
interval = options.interval, // in milliseconds
maxInInterval = options.maxInInterval,
minDifference = options.minDifference || 0,
namespace = options.namespace || "redis-rate-limiter-" + Math.random().toString(36).slice(2);
var result, sha, load;
assert(redis, "`options.redis` must be a redis client");
assert(interval === Infinity || interval > 0 && isInt(options.interval), "`options.interval` must be a positive integer");
assert(maxInInterval > 0 && isInt(maxInInterval), "`options.maxInInterval must be a positive integer");
/*
convert interval and maxInInterval to a fill rate
fillRate = x tokens / millisecond
= maxInInterval / interval
*/
var fillRate = maxInInterval / interval;
/*
capacity = x tokens
= maxInInterval
*/
function redisEvalsha(args, cb) {
redis.evalsha(args, function(err, res) {
if (err) return cb(err);
res = Number(res);
if (res < 0) {
// this is the time until minDifference is satisfied in milliseconds
return cb(null, -Math.floor(res));
} else if (res === 0) {
// no tokens left so the client will have to wait 1 / fillRate milliseconds
return cb(null, 1 / fillRate);
} else {
// there are tokens
return cb(null, 0);
}
});
}
result = function (id, cb) {
if (!cb) {
cb = id;
id = "";
}
assert.equal(typeof cb, "function", "Callback must be a function.");
var now = Date.now();
var key = namespace + id;
var tokenKey = key + ":token";
var lastAccessedKey = key + ":timestamp";
var args = [
sha,
2, // number of keys
tokenKey,
lastAccessedKey,
fillRate,
maxInInterval, // bucket capacity
now,
minDifference || 0,
isFinite(interval) ? Math.ceil(interval / 1000) : 0 // interval in seconds for redis ttl
];
if (!sha) {
load(function(err) {
if (err) return cb(err);
args[0] = sha;
redisEvalsha(args, cb);
})
} else {
redisEvalsha(args, cb);
}
}
result.load = load = function(cb) {
redis.script(["LOAD", lua], function(err, result) {
if (err) return cb(err);
sha = result;
cb(null);
});
}
return result;
}
module.exports = RedisRateLimiter;