prosody
9a35ae935390
plugins/mod_limits.lua

Find changesets by keywords by author, files, the commit message, revision number or hash, or revset expression.

Software / code / prosody

Annotate

plugins/mod_limits.lua @ 9727:9a35ae935390

util.table: Add test for create()
author Kim Alvefur <zash@zash.se>
date Sun, 23 Dec 2018 15:01:37 +0100
parent 8803:60e113f3682f
child 9941:a2f8d54dd445
child 11550:929de6ade6b6
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
8453
6b3e7fddd723 mod_limits: Fix typo in comment
Kim Alvefur <zash@zash.se>
parents: 8269
diff changeset
1 -- Because we deal with pre-authed sessions and streams we can't be host-specific
8256
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";
8269
25237002aba4 mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba)
Matthew Wild <mwild1@gmail.com>
parents: 8256
diff changeset
7 local ceil = math.ceil;
8256
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);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
34 if not n_burst then
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 } }
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
42 local limits = {};
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
43
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
44 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
45 limits[sess_type] = {
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
46 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
47 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
48 };
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
49 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
50
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
51 local default_filter_set = {};
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 function default_filter_set.bytes_in(bytes, session)
8803
60e113f3682f mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents: 8453
diff changeset
54 local sess_throttle = session.throttle;
60e113f3682f mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents: 8453
diff changeset
55 if sess_throttle then
60e113f3682f mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents: 8453
diff changeset
56 local ok, balance, outstanding = sess_throttle:poll(#bytes, true);
8256
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
57 if not ok then
8803
60e113f3682f mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents: 8453
diff changeset
58 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
8269
25237002aba4 mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba)
Matthew Wild <mwild1@gmail.com>
parents: 8256
diff changeset
59 outstanding = ceil(outstanding);
8256
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
60 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
61 local outstanding_data = bytes:sub(-outstanding);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
62 bytes = bytes:sub(1, #bytes-outstanding);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
63 timer.add_task(limits_resolution, function ()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
64 if not session.conn then return; end
8803
60e113f3682f mod_limits: Rename variable to avoid name clash [luacheck]
Kim Alvefur <zash@zash.se>
parents: 8453
diff changeset
65 if sess_throttle:peek(#outstanding_data) then
8256
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
66 session.log("debug", "Resuming paused session");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
67 session.conn:resume();
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
68 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
69 -- 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
70 session.data(outstanding_data);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
71 end);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
72 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
73 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
74 return bytes;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
75 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
76
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
77 local type_filters = {
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
78 c2s = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
79 s2sin = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
80 s2sout = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
81 };
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
82
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
83 local function filter_hook(session)
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
84 local session_type = session.type:match("^[^_]+");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
85 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
86 if opts then
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
87 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
88 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
89 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
90 end
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 function module.load()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
93 filters.add_filter_hook(filter_hook);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
94 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
95
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
96 function module.unload()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
97 filters.remove_filter_hook(filter_hook);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
98 end