-
Notifications
You must be signed in to change notification settings - Fork 150
/
conn.lua
125 lines (88 loc) · 2.24 KB
/
conn.lua
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
-- Copyright (C) Yichun Zhang (agentzh)
--
-- This library is an enhanced Lua port of the standard ngx_limit_conn
-- module.
local math = require "math"
local setmetatable = setmetatable
local floor = math.floor
local ngx_shared = ngx.shared
local assert = assert
local _M = {
_VERSION = '0.08'
}
local mt = {
__index = _M
}
function _M.new(dict_name, max, burst, default_conn_delay)
local dict = ngx_shared[dict_name]
if not dict then
return nil, "shared dict not found"
end
assert(max > 0 and burst >= 0 and default_conn_delay > 0)
local self = {
dict = dict,
max = max + 0, -- just to ensure the param is good
burst = burst,
unit_delay = default_conn_delay,
}
return setmetatable(self, mt)
end
function _M.incoming(self, key, commit)
local dict = self.dict
local max = self.max
self.committed = false
local conn, err
if commit then
conn, err = dict:incr(key, 1, 0)
if not conn then
return nil, err
end
if conn > max + self.burst then
conn, err = dict:incr(key, -1)
if not conn then
return nil, err
end
return nil, "rejected"
end
self.committed = true
else
conn = (dict:get(key) or 0) + 1
if conn > max + self.burst then
return nil, "rejected"
end
end
if conn > max then
-- make the excessive connections wait
return self.unit_delay * floor((conn - 1) / max), conn
end
-- we return a 0 delay by default
return 0, conn
end
function _M.is_committed(self)
return self.committed
end
function _M.leaving(self, key, req_latency)
assert(key)
local dict = self.dict
local conn, err = dict:incr(key, -1)
if not conn then
return nil, err
end
if req_latency then
local unit_delay = self.unit_delay
self.unit_delay = (req_latency + unit_delay) / 2
end
return conn
end
function _M.uncommit(self, key)
assert(key)
local dict = self.dict
return dict:incr(key, -1)
end
function _M.set_conn(self, conn)
self.max = conn
end
function _M.set_burst(self, burst)
self.burst = burst
end
return _M