1 |
|
-module(mongoose_shaper). |
2 |
|
|
3 |
|
-export([child_spec/0]). |
4 |
|
-export([new/1, update/2, wait/5, reset_all_shapers/1]). |
5 |
|
-ignore_xref([reset_all_shapers/1]). |
6 |
|
|
7 |
|
-type shaper() :: opuntia:shaper(). |
8 |
|
-export_type([shaper/0]). |
9 |
|
|
10 |
|
-spec child_spec() -> supervisor:child_spec(). |
11 |
|
child_spec() -> |
12 |
42 |
WPoolOpts = [{workers, 10}, {worker, {opuntia_srv, {?MODULE, #{}}}}], |
13 |
42 |
#{id => ?MODULE, |
14 |
|
start => {wpool, start_pool, [?MODULE, WPoolOpts]}, |
15 |
|
restart => permanent, |
16 |
|
shutdown => infinity, |
17 |
|
type => supervisor}. |
18 |
|
|
19 |
|
-spec new(atom()) -> opuntia:shaper(). |
20 |
|
new(Name) -> |
21 |
12168 |
opuntia:new(get_shaper_config(Name)). |
22 |
|
|
23 |
|
-spec update(opuntia:shaper(), opuntia:tokens()) -> {opuntia:shaper(), opuntia:delay()}. |
24 |
|
update(Shaper, Tokens) -> |
25 |
51968 |
opuntia:update(Shaper, Tokens). |
26 |
|
|
27 |
|
%% @doc Shapes the caller from executing the action. |
28 |
|
-spec wait(HostType :: mongooseim:host_type_or_global(), |
29 |
|
Domain :: jid:lserver(), |
30 |
|
Action :: atom(), |
31 |
|
FromJID :: jid:jid(), |
32 |
|
Size :: integer()) -> continue | {error, max_delay_reached}. |
33 |
|
wait(HostType, Domain, Action, FromJID, Size) -> |
34 |
3948 |
Worker = wpool_pool:hash_worker(?MODULE, FromJID), |
35 |
3948 |
Config = get_shaper_config(get_shaper_name(HostType, Domain, Action, FromJID)), |
36 |
3948 |
Key = new_key(Domain, Action, FromJID), |
37 |
3948 |
opuntia_srv:wait(Worker, Key, Size, Config). |
38 |
|
|
39 |
|
-type key() :: {global | jid:server(), atom(), jid:jid()}. |
40 |
|
-spec new_key(jid:server() | global, atom(), jid:jid()) -> key(). |
41 |
|
new_key(Domain, Action, FromJID) -> |
42 |
3948 |
{Domain, Action, FromJID}. |
43 |
|
|
44 |
|
%% @doc Ask all shaper servers to forget current shapers and read settings again |
45 |
|
reset_all_shapers(_HostType) -> |
46 |
1 |
[ opuntia_srv:reset_shapers(ProcName) || ProcName <- wpool:get_workers(?MODULE) ], |
47 |
1 |
ok. |
48 |
|
|
49 |
|
-spec get_shaper_name(HostType :: mongooseim:host_type_or_global(), |
50 |
|
Domain :: global | jid:server(), |
51 |
|
Action :: atom(), |
52 |
|
FromJID :: jid:jid()) -> allow | none. |
53 |
|
get_shaper_name(HostType, Domain, Action, FromJID) -> |
54 |
3948 |
case acl:match_rule(HostType, Domain, Action, FromJID) of |
55 |
3010 |
deny -> default_shaper(); |
56 |
938 |
Value -> Value |
57 |
|
end. |
58 |
|
|
59 |
|
-spec get_shaper_config(atom()) -> number(). |
60 |
|
get_shaper_config(Name) -> |
61 |
16116 |
case mongoose_config:lookup_opt([shaper, Name]) of |
62 |
|
{ok, #{max_rate := MaxRatePerSecond}} -> |
63 |
938 |
#{bucket_size => MaxRatePerSecond, rate => MaxRatePerSecond, |
64 |
|
time_unit => second, start_full => true}; |
65 |
|
{error, not_found} -> |
66 |
15178 |
0 |
67 |
|
end. |
68 |
|
|
69 |
|
default_shaper() -> |
70 |
3010 |
none. |