%% @private %% @see amoc_throttle %% @copyright 2023 Erlang Solutions Ltd. -module(amoc_throttle_controller). -behaviour(gen_server). %% API -export([start_link/0, ensure_throttle_processes_started/4, pause/1, resume/1, stop/1, change_rate/3, change_rate_gradually/6, run/2, telemetry_event/2]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2]). -define(SERVER, ?MODULE). -define(MASTER_SERVER, {?SERVER, amoc_cluster:master_node()}). -record(throttle_info, { rate :: amoc_throttle:rate(), interval :: amoc_throttle:interval(), no_of_procs :: pos_integer(), active :: boolean(), change_plan :: change_rate_plan() | undefined }). -record(change_rate_plan, { high_rate :: pos_integer(), no_of_steps :: non_neg_integer(), timer :: timer:tref()}). -type name() :: amoc_throttle:name(). -type change_rate_plan() :: #change_rate_plan{}. -type throttle_info() :: #throttle_info{}. -type state() :: #{name() => throttle_info()}. %%%=================================================================== %%% API %%%=================================================================== -spec(start_link() -> {ok, Pid :: pid()} | ignore | {error, Reason :: term()}). start_link() -> pg:start_link(), gen_server:start_link({local, ?SERVER}, ?MODULE, [], []). -spec(ensure_throttle_processes_started(name(), amoc_throttle:rate(), amoc_throttle:interval(), pos_integer()) -> {ok, started_throttle_processes} | {ok, throttle_processes_already_started} | {error, any()}). ensure_throttle_processes_started(Name, Rate, Interval, NoOfProcesses) -> maybe_raise_event(Name, init), gen_server:call(?MASTER_SERVER, {start_processes, Name, Rate, Interval, NoOfProcesses}). -spec run(name(), fun(() -> any())) -> ok | {error, any()}. run(Name, Fn) -> case get_throttle_process(Name) of {ok, Pid} -> maybe_raise_event(Name, request), Fun = fun() -> maybe_raise_event(Name, execute), Fn() end, amoc_throttle_process:run(Pid, Fun), ok; Error -> Error end. -spec pause(name()) -> ok | {error, any()}. pause(Name) -> gen_server:call(?MASTER_SERVER, {pause, Name}). -spec resume(name()) -> ok | {error, any()}. resume(Name) -> gen_server:call(?MASTER_SERVER, {resume, Name}). -spec change_rate(name(), amoc_throttle:rate(), amoc_throttle:interval()) -> ok | {error, any()}. change_rate(Name, Rate, Interval) -> gen_server:call(?MASTER_SERVER, {change_rate, Name, Rate, Interval}). -spec change_rate_gradually(name(), amoc_throttle:rate(), amoc_throttle:rate(), amoc_throttle:interval(), pos_integer(), pos_integer()) -> ok | {error, any()}. change_rate_gradually(Name, LowRate, HighRate, RateInterval, StepInterval, NoOfSteps) -> gen_server:call(?MASTER_SERVER, {change_rate_gradually, Name, LowRate, HighRate, RateInterval, StepInterval, NoOfSteps}). -spec stop(name()) -> ok | {error, any()}. stop(Name) -> gen_server:call(?MASTER_SERVER, {stop, Name}). -spec telemetry_event(name(), request | execute) -> ok. telemetry_event(Name, Event) when Event =:= request; Event =:= execute -> raise_event(Name, Event). %%%=================================================================== %%% gen_server callbacks %%%=================================================================== -spec init([]) -> {ok, #{}}. init([]) -> {ok, #{}}. -spec handle_call({start_processes, name(), pos_integer(), amoc_throttle:interval(), pos_integer()}, From :: {pid(), Tag :: term()}, state()) -> {reply, {ok, started}, state()} | {reply, {error, wrong_no_of_procs}, state()}; ({pause | resume | stop}, From :: {pid(), Tag :: term()}, state()) -> {reply, ok, state()} | {reply, Error :: any(), state()}; ({change_rate, name(), amoc_throttle:rate(), amoc_throttle:interval()}, From :: {pid(), Tag :: term()}, state()) -> {reply, ok, state()} | {reply, {error, any()}, state()}; ({change_rate_gradually, name(), amoc_throttle:rate(), amoc_throttle:rate(), amoc_throttle:interval(), pos_integer(), pos_integer()}, From :: {pid(), Tag :: term()}, state()) -> {reply, ok, state()} | {reply, {error, any()}, state()}. handle_call({start_processes, Name, Rate, Interval, NoOfProcesses}, _From, State) -> case pg:get_members(Name) of [] -> RealNoOfProcesses = start_processes(Name, Rate, Interval, NoOfProcesses), {reply, {ok, started}, State#{Name => #throttle_info{rate = Rate, interval = Interval, active = true, no_of_procs = RealNoOfProcesses}}}; Group when is_list(Group) -> ExpectedNoOfProcesses = min(Rate, NoOfProcesses), case length(Group) of ExpectedNoOfProcesses -> {reply, {ok, started}, State}; _ -> {reply, {error, wrong_no_of_procs}, State} end end; handle_call({pause, Name}, _From, State) -> case all_processes(Name, pause) of ok -> Info = maps:get(Name, State), {reply, ok, State#{Name => Info#throttle_info{active = false}}}; Error -> {reply, Error, State} end; handle_call({resume, Name}, _From, State) -> case all_processes(Name, resume) of ok -> Info = maps:get(Name, State), {reply, ok, State#{Name => Info#throttle_info{active = true}}}; Error -> {reply, Error, State} end; handle_call({change_rate, Name, Rate, Interval}, _From, State) -> Info = maps:get(Name, State), case maybe_change_rate(Name, Rate, Interval, Info) of {ok, Rate} -> UpdatedInfo = Info#throttle_info{rate = Rate, interval = Interval}, {reply, ok, State#{Name => UpdatedInfo}}; Error -> {reply, Error, State} end; handle_call({change_rate_gradually, Name, LowRate, HighRate, RateInterval, StepInterval, NoOfSteps}, _From, State) -> Info = maps:get(Name, State), case Info#throttle_info.change_plan of undefined -> NewInfo = start_gradual_rate_change(Name, LowRate, HighRate, RateInterval, StepInterval, NoOfSteps, Info), {reply, ok, State#{Name => NewInfo}}; _ -> {reply, {error, cannot_change_rate}, State} end; handle_call({stop, Name}, _From, State) -> case all_processes(Name, stop) of ok -> {reply, ok, maps:remove(Name, State)}; Error -> {reply, Error, State} end. -spec(handle_cast(any(), state()) -> {noreply, state()}). handle_cast(_, State) -> {noreply, State}. -spec(handle_info({change_plan, name()}, state()) -> {noreply, state()}). handle_info({change_plan, Name}, State) -> Info = maps:get(Name, State), Plan = Info#throttle_info.change_plan, case Plan#change_rate_plan.no_of_steps of 1 -> NewState = change_rate_and_stop_plan(Name, State), {noreply, NewState}; N when N > 1 -> NewState = continue_plan(Name, State), {noreply, NewState} end. %%%=================================================================== %%% Internal functions %%%=================================================================== maybe_raise_event(Name, Event) -> case amoc_cluster:master_node() =:= node() of true -> ok; _ -> raise_event(Name, Event) end. raise_event(Name, Event) when Event =:= request; Event =:= execute; Event =:= init -> amoc_telemetry:execute([throttle, Event], #{count => 1}, #{name => Name}). report_rate(Name, RatePerMinute) -> amoc_telemetry:execute([throttle, rate], #{rate => RatePerMinute}, #{name => Name}). -spec change_rate_and_stop_plan(name(), state()) -> state(). change_rate_and_stop_plan(Name, State) -> Info = maps:get(Name, State), Plan = Info#throttle_info.change_plan, Interval = Info#throttle_info.interval, TRef = Plan#change_rate_plan.timer, HighRate = Plan#change_rate_plan.high_rate, {ok, HighRate} = do_change_rate(Name, HighRate, Interval), {ok, cancel} = timer:cancel(TRef), consume_all_timer_ticks({change_plan, Name}), State#{Name => Info#throttle_info{rate = HighRate, change_plan = undefined}}. consume_all_timer_ticks(Msg) -> receive Msg -> consume_all_timer_ticks(Msg) after 0 -> ok end. -spec continue_plan(name(), state()) -> state(). continue_plan(Name, State) -> Info = maps:get(Name, State), Plan = Info#throttle_info.change_plan, LowRate = Info#throttle_info.rate, HighRate = Plan#change_rate_plan.high_rate, NoOfSteps = Plan#change_rate_plan.no_of_steps, Step = (HighRate - LowRate) div (NoOfSteps), NewRate = LowRate + Step, {ok, NewRate} = do_change_rate(Name, NewRate, Info#throttle_info.interval), NewPlan = Plan#change_rate_plan{no_of_steps = NoOfSteps - 1}, State#{Name => Info#throttle_info{rate = NewRate, change_plan = NewPlan}}. -spec rate_per_minute(amoc_throttle:rate(), amoc_throttle:interval()) -> amoc_throttle:rate(). rate_per_minute(_, 0) -> 0; rate_per_minute(Rate, Interval) -> (Rate * 60000) div Interval. -spec start_processes(name(), amoc_throttle:rate(), amoc_throttle:interval(), pos_integer()) -> pos_integer(). start_processes(Name, Rate, Interval, NoOfProcesses) -> raise_event(Name, init), RatePerMinute = rate_per_minute(Rate, Interval), report_rate(Name, RatePerMinute), RealNoOfProcesses = min(Rate, NoOfProcesses), start_throttle_processes(Name, Interval, Rate, RealNoOfProcesses), RealNoOfProcesses. -spec get_throttle_process(name()) -> {error, {no_throttle_process_registered, name()}} | {error, any()} | {ok, pid()}. get_throttle_process(Name) -> case pg:get_members(Name) of [] -> {error, {no_throttle_process_registered, Name}}; List -> %% nonempty list N = rand:uniform(length(List)), {ok, lists:nth(N, List)} end. -spec maybe_change_rate(name(), amoc_throttle:rate(), amoc_throttle:interval(), throttle_info()) -> {ok, non_neg_integer()} | {error, any()}. maybe_change_rate(Name, Rate, Interval, Info) -> CurrentRatePerMin = rate_per_minute(Info#throttle_info.rate, Info#throttle_info.interval), ReqRatePerMin = rate_per_minute(Rate, Interval), case {CurrentRatePerMin, Info#throttle_info.change_plan} of {ReqRatePerMin, _} -> {ok, ReqRatePerMin}; {_, undefined} -> do_change_rate(Name, Rate, Interval); _ -> {error, cannot_change_rate} end. -spec do_change_rate(name(), amoc_throttle:rate(), amoc_throttle:interval()) -> {ok, non_neg_integer()} | {error, any()}. do_change_rate(Name, Rate, Interval) -> case pg:get_members(Name) of [] -> {error, no_processes_in_group}; List when is_list(List) -> RatePerMinute = rate_per_minute(Rate, Interval), report_rate(Name, RatePerMinute), update_throttle_processes(List, Interval, Rate, length(List)), {ok, RatePerMinute} end. -spec start_gradual_rate_change( name(), amoc_throttle:rate(), amoc_throttle:rate(), amoc_throttle:interval(), pos_integer(), pos_integer(), throttle_info()) -> throttle_info(). start_gradual_rate_change(Name, LowRate, HighRate, RateInterval, StepInterval, NoOfSteps, Info) -> {ok, LowRate} = do_change_rate(Name, LowRate, RateInterval), {ok, Timer} = timer:send_interval(StepInterval, {change_plan, Name}), Plan = #change_rate_plan{high_rate = HighRate, no_of_steps = NoOfSteps, timer = Timer}, Info#throttle_info{rate = LowRate, interval = RateInterval, change_plan = Plan}. start_throttle_processes(Name, Interval, Rate, 1) -> start_throttle_process(Name, Interval, Rate); start_throttle_processes(Name, Interval, Rate, N) when is_integer(N), N > 1 -> ProcessRate = Rate div N, start_throttle_process(Name, Interval, ProcessRate), start_throttle_processes(Name, Interval, Rate - ProcessRate, N - 1). start_throttle_process(Name, Interval, Rate) -> {ok, Pid} = amoc_throttle_process:start(Name, Interval, Rate), pg:join(Name, Pid). update_throttle_processes([Pid], Interval, Rate, 1) -> amoc_throttle_process:update(Pid, Interval, Rate); update_throttle_processes([Pid | Tail], Interval, Rate, N) when N > 1 -> ProcessRate = Rate div N, amoc_throttle_process:update(Pid, Interval, ProcessRate), update_throttle_processes(Tail, Interval, Rate - ProcessRate, N - 1). all_processes(Name, Cmd) -> case pg:get_members(Name) of [] -> {error, no_processes_in_group}; Ps -> [run_cmd(P, Cmd) || P <- Ps], ok end. run_cmd(Pid, stop) -> amoc_throttle_process:stop(Pid); run_cmd(Pid, pause) -> amoc_throttle_process:pause(Pid); run_cmd(Pid, resume) -> amoc_throttle_process:resume(Pid).