| 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 |
53 |
WPoolOpts = [{workers, 10}, {worker, {opuntia_srv, {?MODULE, #{}}}}], |
| 13 |
53 |
#{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 |
13137 |
opuntia:new(get_shaper_config(Name)). |
| 22 |
|
|
| 23 |
|
-spec update(opuntia:shaper(), opuntia:tokens()) -> {opuntia:shaper(), opuntia:delay()}. |
| 24 |
|
update(Shaper, Tokens) -> |
| 25 |
60023 |
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 |
3936 |
Worker = wpool_pool:hash_worker(?MODULE, FromJID), |
| 35 |
3936 |
Config = get_shaper_config(get_shaper_name(HostType, Domain, Action, FromJID)), |
| 36 |
3936 |
Key = new_key(Domain, Action, FromJID), |
| 37 |
3936 |
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 |
3936 |
{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 |
3936 |
case acl:match_rule(HostType, Domain, Action, FromJID) of |
| 55 |
2998 |
deny -> default_shaper(); |
| 56 |
938 |
Value -> Value |
| 57 |
|
end. |
| 58 |
|
|
| 59 |
|
-spec get_shaper_config(atom()) -> number(). |
| 60 |
|
get_shaper_config(Name) -> |
| 61 |
17073 |
case mongoose_config:lookup_opt([shaper, Name]) of |
| 62 |
|
{ok, #{max_rate := MaxRatePerSecond}} -> |
| 63 |
995 |
#{bucket_size => MaxRatePerSecond, rate => MaxRatePerSecond, |
| 64 |
|
time_unit => second, start_full => true}; |
| 65 |
|
{error, not_found} -> |
| 66 |
16078 |
0 |
| 67 |
|
end. |
| 68 |
|
|
| 69 |
|
default_shaper() -> |
| 70 |
2998 |
none. |