tokenBucket.js
5.43 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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
/**
* A hierarchical token bucket for rate limiting. See
* http://en.wikipedia.org/wiki/Token_bucket for more information.
* @author John Hurliman <jhurliman@cull.tv>
*
* @param {Number} bucketSize Maximum number of tokens to hold in the bucket.
* Also known as the burst rate.
* @param {Number} tokensPerInterval Number of tokens to drip into the bucket
* 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 {TokenBucket} parentBucket Optional. A token bucket that will act as
* the parent of this bucket.
*/
var TokenBucket = function(bucketSize, tokensPerInterval, interval, parentBucket) {
this.bucketSize = bucketSize;
this.tokensPerInterval = tokensPerInterval;
if (typeof interval === 'string') {
switch (interval) {
case 'sec': case 'second':
this.interval = 1000; break;
case 'min': case 'minute':
this.interval = 1000 * 60; break;
case 'hr': case 'hour':
this.interval = 1000 * 60 * 60; break;
case 'day':
this.interval = 1000 * 60 * 60 * 24; break;
default:
throw new Error('Invaid interval ' + interval);
}
} else {
this.interval = interval;
}
this.parentBucket = parentBucket;
this.content = 0;
this.lastDrip = +new Date();
};
TokenBucket.prototype = {
bucketSize: 1,
tokensPerInterval: 1,
interval: 1000,
parentBucket: null,
content: 0,
lastDrip: 0,
/**
* Remove the requested number of tokens and fire the given callback. If the
* bucket (and any parent buckets) contains enough tokens 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) {
var self = this;
// Is this an infinite size bucket?
if (!this.bucketSize) {
process.nextTick(callback.bind(null, null, count, Number.POSITIVE_INFINITY));
return true;
}
// Make sure the bucket can hold the requested number of tokens
if (count > this.bucketSize) {
process.nextTick(callback.bind(null, 'Requested tokens ' + count + ' exceeds bucket size ' +
this.bucketSize, null));
return false;
}
// Drip new tokens into this bucket
this.drip();
// If we don't have enough tokens in this bucket, come back later
if (count > this.content)
return comeBackLater();
if (this.parentBucket) {
// Remove the requested from the parent bucket first
return this.parentBucket.removeTokens(count, function(err, remainingTokens) {
if (err) return callback(err, null);
// Check that we still have enough tokens in this bucket
if (count > self.content)
return comeBackLater();
// Tokens were removed from the parent bucket, now remove them from
// this bucket and fire the callback. Note that we look at the current
// bucket and parent bucket's remaining tokens and return the smaller
// of the two values
self.content -= count;
callback(null, Math.min(remainingTokens, self.content));
});
} else {
// Remove the requested tokens from this bucket and fire the callback
this.content -= count;
process.nextTick(callback.bind(null, null, this.content));
return true;
}
function comeBackLater() {
// How long do we need to wait to make up the difference in tokens?
var waitInterval = Math.ceil(
(count - self.content) * (self.interval / self.tokensPerInterval));
setTimeout(function() { self.removeTokens(count, callback); }, waitInterval);
return false;
}
},
/**
* Attempt to remove the requested number of tokens and return immediately.
* If the bucket (and any parent buckets) contains enough tokens 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) {
// Is this an infinite size bucket?
if (!this.bucketSize)
return true;
// Make sure the bucket can hold the requested number of tokens
if (count > this.bucketSize)
return false;
// Drip new tokens into this bucket
this.drip();
// If we don't have enough tokens in this bucket, return false
if (count > this.content)
return false;
// Try to remove the requested tokens from the parent bucket
if (this.parentBucket && !this.parent.tryRemoveTokens(count))
return false;
// Remove the requested tokens from this bucket and return
this.content -= count;
return true;
},
/**
* Add any new tokens to the bucket since the last drip.
* @returns {Boolean} True if new tokens were added, otherwise false.
*/
drip: function() {
if (!this.tokensPerInterval) {
this.content = this.bucketSize;
return;
}
var now = +new Date();
var deltaMS = Math.max(now - this.lastDrip, 0);
this.lastDrip = now;
var dripAmount = deltaMS * (this.tokensPerInterval / this.interval);
this.content = Math.min(this.content + dripAmount, this.bucketSize);
}
};
module.exports = TokenBucket;