taskforce

-module(tf_master_serv).

-ifdef(E48).
-moduledoc false.
-endif.

-behaviour(gen_server).

%% API functions
-export([start_link/2]).

%% gen_server callbacks
-export([
    init/1,
    handle_call/3,
    handle_cast/2,
    handle_info/2,
    terminate/2,
    code_change/3
]).

-ignore_xref([{start_link, 2}]).

-include("taskforce.hrl").

-type exec_state() :: idle | {running, Bidder :: any()}.

-record(master_state, {
    patron_pid :: pid(),
    patron_monitor :: reference(),
    max_minion_count :: pos_integer(),
    exec_state = 'idle' :: exec_state(),
    tasks = [] :: [tf_task()],
    consumed_task_ids = [] :: [any()],
    completed = [] :: [any()],
    timedout_task_ids = [] :: [any()]
}).
-type state() :: #master_state{}.

-if(?OTP_RELEASE =< 24).
-type gen_server_from() :: {pid(), reference()}.
-else.
-type gen_server_from() :: gen_server:from().
-endif.

%%%===================================================================
%%% API functions
%%%===================================================================

-spec start_link(pid(), pos_integer()) -> {ok, pid()}.
start_link(PatronPid, MaxMinionCount) ->
    gen_server:start_link(?MODULE, [PatronPid, MaxMinionCount], []).

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

-spec init([pid() | pos_integer(), ...]) -> {ok, state()}.
init([PatronPid, MaxMinionCount]) ->
    State0 = #master_state{
        patron_pid = PatronPid,
        patron_monitor = monitor(process, PatronPid),
        max_minion_count = MaxMinionCount
    },
    {ok, State0}.

-spec handle_call(term(), gen_server_from(), state()) ->
    {reply, {ok, tf_task()} | {error, no_more_tasks}, state()}
    | {noreply, state()}
    | {stop, normal, state()}.
handle_call(
    {do_my_bidding, #tf_bidding{} = Bidding}, Bidder, #master_state{exec_state = idle} = State
) ->
    #tf_bidding{
        tasks = Tasks,
        timeout = BiddingTimeout
    } = Bidding,

    MaxMinionCount = State#master_state.max_minion_count,
    MinionCount = min(length(Tasks), MaxMinionCount),
    ok = spawn_minions(MinionCount),

    NewExecState = {running, Bidder},
    NewState = State#master_state{
        exec_state = NewExecState,
        tasks = Tasks
    },
    case Tasks of
        [] ->
            {stop, normal, NewState};
        [_ | _] ->
            erlang:send_after(BiddingTimeout, self(), bidding_timeout),
            {noreply, NewState}
    end;
handle_call(
    consume_task,
    {_MinionPid, _},
    #master_state{exec_state = {running, _}, tasks = [Task | Remaining]} = State
) ->
    NewState = State#master_state{
        tasks = Remaining,
        consumed_task_ids = [
            Task#tf_task.id
            | State#master_state.consumed_task_ids
        ]
    },
    {reply, {ok, Task}, NewState};
handle_call(
    consume_task, {_MinionPid, _}, #master_state{exec_state = {running, _}, tasks = []} = State
) ->
    {reply, {error, no_more_tasks}, State}.

-spec handle_cast(term(), state()) -> {noreply, state()} | {stop, normal, state()}.
handle_cast(
    {{task_completed, TaskId}, TaskResult}, #master_state{exec_state = {running, _}} = State
) ->
    NewCompleted = [{TaskId, TaskResult} | State#master_state.completed],
    NewState = State#master_state{completed = NewCompleted},
    stop_when_done(NewState);
handle_cast({{task_timeout, TaskId}}, #master_state{exec_state = {running, _}} = State) ->
    NewTimeouts = [TaskId | State#master_state.timedout_task_ids],
    NewState = State#master_state{timedout_task_ids = NewTimeouts},
    stop_when_done(NewState).

-spec handle_info(term(), state()) ->
    {stop, {shutdown, patron_death | bidding_timeout}, state()}.
handle_info(
    {'DOWN', Reference, process, _Pid, _Reason}, #master_state{patron_monitor = Reference} = State
) ->
    {stop, {shutdown, patron_death}, State};
handle_info(bidding_timeout, #master_state{exec_state = {running, _}} = State) ->
    {stop, {shutdown, bidding_timeout}, State}.

-spec terminate(term(), state()) -> ok.
terminate(_Reason, #master_state{exec_state = idle} = _State) ->
    ok;
terminate(_Reason, #master_state{exec_state = {running, Bidder}} = State) ->
    #master_state{
        tasks = NeverConsumed,
        consumed_task_ids = ConsumedIds,
        completed = Completed
    } = State,

    NeverConsumedIds = [Task#tf_task.id || Task <- NeverConsumed],
    CompletedIds = [TaskId || {TaskId, _} <- Completed],
    IndividualTimeouts = (ConsumedIds -- CompletedIds),
    GlobalTimeouts = NeverConsumedIds,

    BiddingResults = #tf_bidding_result{
        completed = Completed,
        individual_timeouts = IndividualTimeouts,
        global_timeouts = GlobalTimeouts
    },
    gen_server:reply(Bidder, {ok, BiddingResults}),
    ok.

-spec code_change(term(), state(), term()) -> {ok, state()}.
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%%===================================================================
%%% Internal functions
%%%===================================================================

-spec stop_when_done(state()) -> {noreply, state()} | {stop, normal, state()}.
stop_when_done(State) ->
    case all_tasks_accounted_for(State) of
        true ->
            {stop, normal, State};
        false ->
            {noreply, State}
    end.

-spec all_tasks_accounted_for(state()) -> boolean().
all_tasks_accounted_for(#master_state{
    tasks = Tasks,
    completed = Completed,
    timedout_task_ids = Timeouts,
    consumed_task_ids = ConsumedIds
}) ->
    Tasks =:= [] andalso
        (length(Completed) + length(Timeouts)) =:= length(ConsumedIds).

spawn_minions(MaxMinionCount) ->
    Self = self(),
    lists:foreach(
        fun(_Id) ->
            {ok, _MinionPid} = supervisor:start_child(tf_minion_sup, [Self])
        end,
        lists:seq(1, MaxMinionCount)
    ).