149 lines
4.3 KiB
Erlang
149 lines
4.3 KiB
Erlang
%%%----------------------------------------------------------------------
|
|
%%% File : shaper.erl
|
|
%%% Author : Alexey Shchepin <alexey@process-one.net>
|
|
%%% Purpose : Functions to control connections traffic
|
|
%%% Created : 9 Feb 2003 by Alexey Shchepin <alexey@process-one.net>
|
|
%%%
|
|
%%%
|
|
%%% ejabberd, Copyright (C) 2002-2015 ProcessOne
|
|
%%%
|
|
%%% This program is free software; you can redistribute it and/or
|
|
%%% modify it under the terms of the GNU General Public License as
|
|
%%% published by the Free Software Foundation; either version 2 of the
|
|
%%% License, or (at your option) any later version.
|
|
%%%
|
|
%%% This program is distributed in the hope that it will be useful,
|
|
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
|
|
%%% General Public License for more details.
|
|
%%%
|
|
%%% You should have received a copy of the GNU General Public License along
|
|
%%% with this program; if not, write to the Free Software Foundation, Inc.,
|
|
%%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
|
|
%%%
|
|
%%%----------------------------------------------------------------------
|
|
|
|
-module(shaper).
|
|
|
|
-behaviour(ejabberd_config).
|
|
|
|
-author('alexey@process-one.net').
|
|
|
|
-export([start/0, new/1, new1/1, update/2,
|
|
get_max_rate/1, transform_options/1, load_from_config/0,
|
|
opt_type/1]).
|
|
|
|
-include("ejabberd.hrl").
|
|
-include("logger.hrl").
|
|
|
|
-record(maxrate, {maxrate = 0 :: integer(),
|
|
lastrate = 0.0 :: float(),
|
|
lasttime = 0 :: integer()}).
|
|
|
|
-record(shaper, {name :: {atom(), global},
|
|
maxrate :: integer()}).
|
|
|
|
-type shaper() :: none | #maxrate{}.
|
|
|
|
-export_type([shaper/0]).
|
|
|
|
-spec start() -> ok.
|
|
|
|
start() ->
|
|
mnesia:create_table(shaper,
|
|
[{ram_copies, [node()]},
|
|
{local_content, true},
|
|
{attributes, record_info(fields, shaper)}]),
|
|
mnesia:add_table_copy(shaper, node(), ram_copies),
|
|
load_from_config(),
|
|
ok.
|
|
|
|
-spec load_from_config() -> ok | {error, any()}.
|
|
|
|
load_from_config() ->
|
|
Shapers = ejabberd_config:get_option(
|
|
shaper, fun(V) -> V end, []),
|
|
case mnesia:transaction(
|
|
fun() ->
|
|
lists:foreach(
|
|
fun({Name, MaxRate}) ->
|
|
mnesia:write(#shaper{name = {Name, global},
|
|
maxrate = MaxRate})
|
|
end, Shapers)
|
|
end) of
|
|
{atomic, ok} ->
|
|
ok;
|
|
Err ->
|
|
{error, Err}
|
|
end.
|
|
|
|
-spec get_max_rate(atom()) -> none | non_neg_integer().
|
|
|
|
get_max_rate(none) ->
|
|
none;
|
|
get_max_rate(Name) ->
|
|
case ets:lookup(shaper, {Name, global}) of
|
|
[#shaper{maxrate = R}] ->
|
|
R;
|
|
[] ->
|
|
none
|
|
end.
|
|
|
|
-spec new(atom()) -> shaper().
|
|
|
|
new(none) ->
|
|
none;
|
|
new(Name) ->
|
|
MaxRate = case ets:lookup(shaper, {Name, global}) of
|
|
[#shaper{maxrate = R}] ->
|
|
R;
|
|
[] ->
|
|
none
|
|
end,
|
|
new1(MaxRate).
|
|
|
|
-spec new1(none | integer()) -> shaper().
|
|
|
|
new1(none) -> none;
|
|
new1(MaxRate) ->
|
|
#maxrate{maxrate = MaxRate, lastrate = 0.0,
|
|
lasttime = now_to_usec(now())}.
|
|
|
|
-spec update(shaper(), integer()) -> {shaper(), integer()}.
|
|
|
|
update(none, _Size) -> {none, 0};
|
|
update(#maxrate{} = State, Size) ->
|
|
MinInterv = 1000 * Size /
|
|
(2 * State#maxrate.maxrate - State#maxrate.lastrate),
|
|
Interv = (now_to_usec(now()) - State#maxrate.lasttime) /
|
|
1000,
|
|
?DEBUG("State: ~p, Size=~p~nM=~p, I=~p~n",
|
|
[State, Size, MinInterv, Interv]),
|
|
Pause = if MinInterv > Interv ->
|
|
1 + trunc(MinInterv - Interv);
|
|
true -> 0
|
|
end,
|
|
NextNow = now_to_usec(now()) + Pause * 1000,
|
|
{State#maxrate{lastrate =
|
|
(State#maxrate.lastrate +
|
|
1000000 * Size / (NextNow - State#maxrate.lasttime))
|
|
/ 2,
|
|
lasttime = NextNow},
|
|
Pause}.
|
|
|
|
transform_options(Opts) ->
|
|
lists:foldl(fun transform_options/2, [], Opts).
|
|
|
|
transform_options({OptName, Name, {maxrate, N}}, Opts) when OptName == shaper ->
|
|
[{shaper, [{Name, N}]}|Opts];
|
|
transform_options({OptName, Name, none}, Opts) when OptName == shaper ->
|
|
[{shaper, [{Name, none}]}|Opts];
|
|
transform_options(Opt, Opts) ->
|
|
[Opt|Opts].
|
|
|
|
now_to_usec({MSec, Sec, USec}) ->
|
|
(MSec * 1000000 + Sec) * 1000000 + USec.
|
|
|
|
opt_type(shaper) -> fun (V) -> V end;
|
|
opt_type(_) -> [shaper].
|