author | Kim Alvefur <zash@zash.se> |
Thu, 20 Jan 2022 10:51:46 +0100 | |
branch | 0.11 |
changeset 12206 | ebeb4d959fb3 |
parent 11558 | db8e41eb6eff |
child 11564 | 3bbb1af92514 |
permissions | -rw-r--r-- |
8456
6b3e7fddd723
mod_limits: Fix typo in comment
Kim Alvefur <zash@zash.se>
parents:
8272
diff
changeset
|
1 |
-- Because we deal with pre-authed sessions and streams we can't be host-specific |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
2 |
module:set_global(); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
3 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
4 |
local filters = require "util.filters"; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
5 |
local throttle = require "util.throttle"; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
6 |
local timer = require "util.timer"; |
8272
25237002aba4
mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba)
Matthew Wild <mwild1@gmail.com>
parents:
8259
diff
changeset
|
7 |
local ceil = math.ceil; |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
8 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
9 |
local limits_cfg = module:get_option("limits", {}); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
10 |
local limits_resolution = module:get_option_number("limits_resolution", 1); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
11 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
12 |
local default_bytes_per_second = 3000; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
13 |
local default_burst = 2; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
14 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
15 |
local rate_units = { b = 1, k = 3, m = 6, g = 9, t = 12 } -- Plan for the future. |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
16 |
local function parse_rate(rate, sess_type) |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
17 |
local quantity, unit, exp; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
18 |
if rate then |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
19 |
quantity, unit = rate:match("^(%d+) ?([^/]+)/s$"); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
20 |
exp = quantity and rate_units[unit:sub(1,1):lower()]; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
21 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
22 |
if not exp then |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
23 |
module:log("error", "Error parsing rate for %s: %q, using default rate (%d bytes/s)", sess_type, rate, default_bytes_per_second); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
24 |
return default_bytes_per_second; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
25 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
26 |
return quantity*(10^exp); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
27 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
28 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
29 |
local function parse_burst(burst, sess_type) |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
30 |
if type(burst) == "string" then |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
31 |
burst = burst:match("^(%d+) ?s$"); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
32 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
33 |
local n_burst = tonumber(burst); |
11554
929de6ade6b6
mod_limits: Don't emit error when no burst period is configured
Matthew Wild <mwild1@gmail.com>
parents:
8806
diff
changeset
|
34 |
if burst and not n_burst then |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
35 |
module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, tostring(burst), default_burst); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
36 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
37 |
return n_burst or default_burst; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
38 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
39 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
40 |
-- Process config option into limits table: |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
41 |
-- limits = { c2s = { bytes_per_second = X, burst_seconds = Y } } |
11558
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
42 |
local limits = { |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
43 |
c2s = { |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
44 |
bytes_per_second = 10 * 1024; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
45 |
burst_seconds = 2; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
46 |
}; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
47 |
s2sin = { |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
48 |
bytes_per_second = 30 * 1024; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
49 |
burst_seconds = 2; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
50 |
}; |
db8e41eb6eff
mod_limits: Use default limits if none configured
Matthew Wild <mwild1@gmail.com>
parents:
11554
diff
changeset
|
51 |
}; |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
52 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
53 |
for sess_type, sess_limits in pairs(limits_cfg) do |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
54 |
limits[sess_type] = { |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
55 |
bytes_per_second = parse_rate(sess_limits.rate, sess_type); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
56 |
burst_seconds = parse_burst(sess_limits.burst, sess_type); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
57 |
}; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
58 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
59 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
60 |
local default_filter_set = {}; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
61 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
62 |
function default_filter_set.bytes_in(bytes, session) |
8806
60e113f3682f
mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents:
8456
diff
changeset
|
63 |
local sess_throttle = session.throttle; |
60e113f3682f
mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents:
8456
diff
changeset
|
64 |
if sess_throttle then |
60e113f3682f
mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents:
8456
diff
changeset
|
65 |
local ok, balance, outstanding = sess_throttle:poll(#bytes, true); |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
66 |
if not ok then |
8806
60e113f3682f
mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents:
8456
diff
changeset
|
67 |
session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding); |
8272
25237002aba4
mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba)
Matthew Wild <mwild1@gmail.com>
parents:
8259
diff
changeset
|
68 |
outstanding = ceil(outstanding); |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
69 |
session.conn:pause(); -- Read no more data from the connection until there is no outstanding data |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
70 |
local outstanding_data = bytes:sub(-outstanding); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
71 |
bytes = bytes:sub(1, #bytes-outstanding); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
72 |
timer.add_task(limits_resolution, function () |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
73 |
if not session.conn then return; end |
8806
60e113f3682f
mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents:
8456
diff
changeset
|
74 |
if sess_throttle:peek(#outstanding_data) then |
8259
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
75 |
session.log("debug", "Resuming paused session"); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
76 |
session.conn:resume(); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
77 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
78 |
-- Handle what we can of the outstanding data |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
79 |
session.data(outstanding_data); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
80 |
end); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
81 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
82 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
83 |
return bytes; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
84 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
85 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
86 |
local type_filters = { |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
87 |
c2s = default_filter_set; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
88 |
s2sin = default_filter_set; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
89 |
s2sout = default_filter_set; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
90 |
}; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
91 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
92 |
local function filter_hook(session) |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
93 |
local session_type = session.type:match("^[^_]+"); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
94 |
local filter_set, opts = type_filters[session_type], limits[session_type]; |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
95 |
if opts then |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
96 |
session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
97 |
filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
98 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
99 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
100 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
101 |
function module.load() |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
102 |
filters.add_filter_hook(filter_hook); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
103 |
end |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
104 |
|
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
105 |
function module.unload() |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
106 |
filters.remove_filter_hook(filter_hook); |
cdffe33efae4
mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff
changeset
|
107 |
end |