-
Notifications
You must be signed in to change notification settings - Fork 1
/
index.js
120 lines (110 loc) · 2.9 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
/**
* Created by koala on 2016/9/14.
*/
'use strict';
module.exports = RateLimit;
/**
* 初始化一块用户内存
* id: 单一用户
* currentTime: 当前访问的时间戳
* currentCount: 当前剩余可访问次数,大于0表示可以访问,否则应该拒绝
*/
function createObject(interfaceId, id, currentTime, currentCount) {
return {
interfaceId,
id,
currentTime,
currentCount
}
}
/**
* 这是一个令牌桶对象
* @param options
* @constructor
*/
function RateLimit(options) {
options = options || {};
this.id = options.id || 666;
this.db = options.db;
this.interfaceId = options.interfaceId;
this.limit = options.limit || 100;
this.interval = options.interval || 100000;
this.prefix = 'rate_limit:' + this.interfaceId + ':' + this.id;
this.initBucket((err, obj) => {
if (typeof obj === 'string')
obj = JSON.parse(obj);
this.obj = obj;
});
}
/**
* 检查是否可以访问此接口
* @param callback
*/
RateLimit.prototype.check = function (callback) {
// TUDO
let can = this.addToken();
if (!can) {
return callback('too fast');
}
this.getObj(callback);
};
/**
* 判断是否需要增加令牌
* @returns {boolean}
*/
RateLimit.prototype.addToken = function () {
let now = Date.now();
// 即将要增加的令牌数
let addedTokenCount = Math.floor((now - this.obj.currentTime) * this.limit / this.interval);
this.obj.currentCount = this.obj.currentCount + addedTokenCount;
// 不能超过最大值
this.obj.currentCount = this.obj.currentCount > this.limit ? this.limit : this.obj.currentCount;
if (addedTokenCount) {
this.obj.currentTime = now;
}
if (this.obj.currentCount <= 0) {
// this.obj.currentCount = 0;
// this.setKey(this.prefix, JSON.stringify(this.obj));
return false;
}
this.obj.currentCount--;
this.setKey(this.prefix, JSON.stringify(this.obj), this.interval / 1000);
return true;
};
/**
* 分配一个独立空间给单一用户(redis 的 hset)
*/
RateLimit.prototype.initBucket = function (callback) {
this.getObj((err, data) => {
if (!data) {
data = createObject(this.interfaceId, this.id, Date.now(), this.limit / 2);
this.setKey(this.prefix, JSON.stringify(data), this.interval / 1000);
}
callback(null, data);
});
};
/**
* 从redis获取令牌桶
* @param callback
*/
RateLimit.prototype.getObj = function (callback) {
this.db.get(this.prefix, (err, reply) => {
if (err) {
callback(err);
} else {
callback(null, reply);
}
})
};
/**
* 写入一个key值
* @param key
* @param value
* @param expire
*/
RateLimit.prototype.setKey = function (key, value, expire) {
this.db.set(key, value);
if (expire) {
this.db.expire(key, expire);
}
};