-module(eharbor_coordinator).
-behavior(gen_server).
-export([start_link/1, pause/1, pause/2, get_stats/1]).
-export([init/1, handle_info/2, handle_cast/2, handle_call/3]).
-define(PAUSE_DEFAULT_MS, 100).
-record(state,
{name,
dedup,
ordered,
backlog,
%% How many unique tasks
piers,
%% Breakwater_limit total number of processes waiting out of buffer
breakwater_limit,
breakwater_usage = 0,
piers_in_use = 0,
active_work_registry = #{},
conductor_registry = #{},
conductor_monitors = #{},
pobox_pid = nil,
notify_ready = true,
total_messages = 0,
dropped_messages = 0}).
%% API functions
-spec start_link(eharbor:incomplete_config()) -> {ok, pid()}.
start_link(IncompleteConfig) ->
Config = #{name := Name} = eharbor:merge_with_defaults(IncompleteConfig),
gen_server:start_link({local, eharbor:coordinator_name(Name)}, ?MODULE, Config, []).
%% For testing only simulate a busy coordinator
pause(Name) ->
pause(Name, ?PAUSE_DEFAULT_MS).
pause(Name, Sleep) ->
gen_server:cast(Name, {pause, Sleep}).
%% Get messages vs dropped message count since the start.
get_stats(Name) ->
gen_server:call(Name, get_stats).
%% gen_server callbacks
init(_Config =
#{name := Name,
dedup := Dedup,
ordered := Ordered,
piers := Piers,
breakwater_limit := BreakwaterLimit,
backlog := Backlog}) ->
%% Start the pobox load shedding buffer
{ok, PoboxPid} =
pobox:start_link({local, eharbor:pobox_name(Name)}, erlang:self(), Backlog, keep_old),
%% Enter notify mode (tells pobox we are ready to handle new messages)
ok = pobox:notify(PoboxPid),
{ok,
#state{name = Name,
dedup = Dedup,
ordered = Ordered,
piers = Piers,
breakwater_limit = BreakwaterLimit,
backlog = Backlog,
pobox_pid = PoboxPid}}.
handle_call(get_stats,
_From,
State = #state{total_messages = TotalMessages, dropped_messages = DroppedMessages}) ->
{reply,
#{total_messages => TotalMessages, dropped_messages => DroppedMessages},
handle_notify(State)};
handle_call(_Msg, _From, State) ->
{reply, ok, handle_notify(State)}.
handle_cast({pause, Sleep}, State) ->
timer:sleep(Sleep),
{noreply, handle_notify(State)};
handle_cast(_Msg, State) ->
{noreply, handle_notify(State)}.
handle_info({mail, PoboxPid, new_data},
State =
#state{pobox_pid = PoboxPid,
piers = Piers,
piers_in_use = PiersInUse,
breakwater_usage = BreakwaterUsage,
breakwater_limit = BreakwaterLimit}) ->
%% New data is available in our buffer determine the maximum we can accept without
%% going over our limits and switch the pobox to active mode to get a batch.
AvailableSlots = erlang:min(Piers - PiersInUse, BreakwaterLimit - BreakwaterUsage),
pobox:active(PoboxPid,
fun (_Msg, Count) when Count =:= AvailableSlots ->
skip;
(Msg, Count) ->
{{ok, Msg}, Count + 1}
end,
0),
{noreply, State};
handle_info({mail, PoboxPid, Messages, Count, DropCount},
State =
#state{pobox_pid = PoboxPid,
piers_in_use = PiersInUse,
breakwater_usage = BreakwaterUsage,
active_work_registry = ActiveWorkRegistry0,
conductor_registry = ConductorRegistry0,
conductor_monitors = ConductorMonitors0,
dedup = Dedup,
total_messages = TotalMessages,
dropped_messages = DroppedMessages}) ->
ArgumentGroups =
case Dedup of
true ->
group_by_key([{Args, From} || {'$habor_dock', From, Args} <- Messages]);
false ->
maps:from_list([{Args, [From]} || {'$habor_dock', From, Args} <- Messages])
end,
ArgumentGroupsKeys = maps:keys(ArgumentGroups),
{OldArgsList, NewArgsList} =
lists:partition(fun(ArgumentGroup) -> maps:is_key(ArgumentGroup, ActiveWorkRegistry0) end,
ArgumentGroupsKeys),
[begin
{ConductorPid, _MRef} =
erlang:hd(
maps:get(OldArgs, ActiveWorkRegistry0)),
FollowerFroms = maps:get(OldArgs, ArgumentGroups),
lists:foreach(fun(FollowerFrom) -> gen:reply(FollowerFrom, {follower, ConductorPid}) end,
FollowerFroms)
end
|| OldArgs <- OldArgsList],
NewConductorMrefs =
lists:foldl(fun(NewArgs, Acc) ->
ConductorFrom = hd(maps:get(NewArgs, ArgumentGroups)),
{ConductorPid, _MRef} = ConductorFrom,
ConductorMRef = erlang:monitor(process, ConductorPid),
Acc#{ConductorPid => {ConductorMRef, ConductorFrom}}
end,
#{},
NewArgsList),
NewConductors =
maps:from_list(
lists:map(fun(NewArgs) ->
ConductorFrom = hd(maps:get(NewArgs, ArgumentGroups)),
{ConductorPid, _Mref} = ConductorFrom,
FollowerFroms = tl(maps:get(NewArgs, ArgumentGroups)),
gen:reply(ConductorFrom, conductor),
lists:foreach(fun(FollowerFrom) ->
gen:reply(FollowerFrom, {follower, ConductorPid})
end,
FollowerFroms),
{ConductorFrom, NewArgs}
end,
NewArgsList)),
ConductorRegistry1 = maps:merge(ConductorRegistry0, NewConductors),
ActiveWorkRegistry1 =
maps:merge_with(fun(_K, V1, V2) -> V1 ++ V2 end, ActiveWorkRegistry0, ArgumentGroups),
ConductorMonitors1 = maps:merge(ConductorMonitors0, NewConductorMrefs),
{noreply,
handle_notify(State#state{active_work_registry = ActiveWorkRegistry1,
conductor_registry = ConductorRegistry1,
conductor_monitors = ConductorMonitors1,
piers_in_use = PiersInUse + erlang:length(NewArgsList),
breakwater_usage = BreakwaterUsage + Count,
notify_ready = true,
total_messages = TotalMessages + Count,
dropped_messages = DroppedMessages + DropCount})};
handle_info({'$harbor_work_completed', ConductorFrom},
State =
#state{active_work_registry = ActiveWorkRegistry0,
conductor_registry = ConductorRegistry0,
conductor_monitors = ConductorMonitors0,
breakwater_usage = BreakwaterUsage,
piers_in_use = PiersInUse}) ->
ConductorArgs = maps:get(ConductorFrom, ConductorRegistry0),
{[ConductorFrom | Followers] = BreakwaterProcs, ActiveWorkRegistry1} =
maps:take(ConductorArgs, ActiveWorkRegistry0),
gen:reply(ConductorFrom, {followers, Followers}),
{ConductorArgs, ConductorRegistry1} = maps:take(ConductorFrom, ConductorRegistry0),
{ConductorPid, _ConductorMRef} = ConductorFrom,
{{ConductorMRef, _ConductorFrom}, ConductorMonitors1} =
maps:take(ConductorPid, ConductorMonitors0),
erlang:demonitor(ConductorMRef, [flush]),
{noreply,
handle_notify(State#state{active_work_registry = ActiveWorkRegistry1,
conductor_registry = ConductorRegistry1,
conductor_monitors = ConductorMonitors1,
breakwater_usage = BreakwaterUsage - erlang:length(BreakwaterProcs),
piers_in_use = PiersInUse - 1})};
handle_info({'DOWN', _ConductorMRef, process, ConductorPid, _Reason},
State =
#state{active_work_registry = ActiveWorkRegistry0,
conductor_registry = ConductorRegistry0,
conductor_monitors = ConductorMonitors0,
breakwater_usage = BreakwaterUsage,
piers_in_use = PiersInUse}) ->
{{ConductorMRef, ConductorFrom}, ConductorMonitors1} =
maps:take(ConductorPid, ConductorMonitors0),
{Args, ConductorRegistry1} = maps:take(ConductorFrom, ConductorRegistry0),
{Froms, ActiveWorkRegistry1} = maps:take(Args, ActiveWorkRegistry0),
erlang:demonitor(ConductorMRef, [flush]),
{noreply,
handle_notify(State#state{active_work_registry = ActiveWorkRegistry1,
conductor_registry = ConductorRegistry1,
conductor_monitors = ConductorMonitors1,
breakwater_usage = BreakwaterUsage - erlang:length(Froms),
piers_in_use = PiersInUse - 1})};
handle_info(ping, State) ->
{noreply, handle_notify(State)};
handle_info(_Msg, State) ->
{noreply, handle_notify(State)}.
handle_notify(State =
#state{pobox_pid = PoboxPid,
breakwater_usage = BreakwaterUsage,
breakwater_limit = BreakwaterLimit,
piers_in_use = PiersInUse,
piers = Piers,
notify_ready = true}) ->
{message_queue_len, QueueLen} =
erlang:process_info(
erlang:self(), message_queue_len),
case QueueLen =:= 0 andalso BreakwaterUsage < BreakwaterLimit andalso PiersInUse < Piers
of
true ->
pobox:notify(PoboxPid),
State#state{notify_ready = false};
false ->
State
end;
handle_notify(State) ->
State.
group_by_key(KVs) ->
lists:foldl(fun({K, V}, Acc) ->
case maps:is_key(K, Acc) of
true ->
Value = maps:get(K, Acc),
maps:put(K, [V | Value], Acc);
false ->
maps:put(K, [V], Acc)
end
end,
#{},
KVs).
