rateLimiter.js
4.8 KB
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
var TokenBucket = require('./tokenBucket');
var getMilliseconds = require('./clock');
/**
* A generic rate limiter. Underneath the hood, this uses a token bucket plus
* an additional check to limit how many tokens we can remove each interval.
* @author John Hurliman <jhurliman@jhurliman.org>
*
* @param {Number} tokensPerInterval Maximum number of tokens that can be
* removed at any given moment and over the course of one interval.
* @param {String|Number} interval The interval length in milliseconds, or as
* one of the following strings: 'second', 'minute', 'hour', day'.
* @param {Boolean} fireImmediately Optional. Whether or not the callback
* will fire immediately when rate limiting is in effect (default is false).
*/
var RateLimiter = function(tokensPerInterval, interval, fireImmediately) {
this.tokenBucket = new TokenBucket(tokensPerInterval, tokensPerInterval,
interval, null);
// Fill the token bucket to start
this.tokenBucket.content = tokensPerInterval;
this.curIntervalStart = getMilliseconds();
this.tokensThisInterval = 0;
this.fireImmediately = fireImmediately;
};
RateLimiter.prototype = {
tokenBucket: null,
curIntervalStart: 0,
tokensThisInterval: 0,
fireImmediately: false,
/**
* Remove the requested number of tokens and fire the given callback. If the
* rate limiter contains enough tokens and we haven't spent too many tokens
* in this interval already, this will happen immediately. Otherwise, the
* removal and callback will happen when enough tokens become available.
* @param {Number} count The number of tokens to remove.
* @param {Function} callback(err, remainingTokens)
* @returns {Boolean} True if the callback was fired immediately, otherwise
* false.
*/
removeTokens: function(count, callback) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize) {
process.nextTick(callback.bind(null, 'Requested tokens ' + count +
' exceeds maximum tokens per interval ' + this.tokenBucket.bucketSize,
null));
return false;
}
var self = this;
var now = getMilliseconds();
// Advance the current interval and reset the current interval token count
// if needed
if (now < this.curIntervalStart
|| now - this.curIntervalStart >= this.tokenBucket.interval) {
this.curIntervalStart = now;
this.tokensThisInterval = 0;
}
// If we don't have enough tokens left in this interval, wait until the
// next interval
if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval) {
if (this.fireImmediately) {
process.nextTick(callback.bind(null, null, -1));
} else {
var waitInterval = Math.ceil(
this.curIntervalStart + this.tokenBucket.interval - now);
setTimeout(function() {
self.tokenBucket.removeTokens(count, afterTokensRemoved);
}, waitInterval);
}
return false;
}
// Remove the requested number of tokens from the token bucket
return this.tokenBucket.removeTokens(count, afterTokensRemoved);
function afterTokensRemoved(err, tokensRemaining) {
if (err) return callback(err, null);
self.tokensThisInterval += count;
callback(null, tokensRemaining);
}
},
/**
* Attempt to remove the requested number of tokens and return immediately.
* If the bucket (and any parent buckets) contains enough tokens and we
* haven't spent too many tokens in this interval already, this will return
* true. Otherwise, false is returned.
* @param {Number} count The number of tokens to remove.
* @param {Boolean} True if the tokens were successfully removed, otherwise
* false.
*/
tryRemoveTokens: function(count) {
// Make sure the request isn't for more than we can handle
if (count > this.tokenBucket.bucketSize)
return false;
var now = getMilliseconds();
// Advance the current interval and reset the current interval token count
// if needed
if (now < this.curIntervalStart
|| now - this.curIntervalStart >= this.tokenBucket.interval) {
this.curIntervalStart = now;
this.tokensThisInterval = 0;
}
// If we don't have enough tokens left in this interval, return false
if (count > this.tokenBucket.tokensPerInterval - this.tokensThisInterval)
return false;
// Try to remove the requested number of tokens from the token bucket
var removed = this.tokenBucket.tryRemoveTokens(count);
if (removed) {
this.tokensThisInterval += count;
}
return removed;
},
/**
* Returns the number of tokens remaining in the TokenBucket.
* @returns {Number} The number of tokens remaining.
*/
getTokensRemaining: function () {
this.tokenBucket.drip();
return this.tokenBucket.content;
}
};
module.exports = RateLimiter;