shigoto

-module(shigoto_queue).
-moduledoc """
Per-queue gen_server that polls for available jobs and dispatches
them to the executor supervisor. Supports graceful shutdown, queue
weights, and pause/resume.
""".
-behaviour(gen_server).

-export([start_link/2, pause/1, resume/1]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2]).

-record(state, {
    queue :: binary(),
    concurrency :: pos_integer(),
    active :: non_neg_integer(),
    pool :: atom(),
    paused :: boolean(),
    shutting_down :: boolean(),
    executors :: #{pid() => reference()},
    weight :: pos_integer(),
    weight_counter :: non_neg_integer(),
    fair :: boolean()
}).

-doc false.
start_link(Queue, Concurrency) ->
    gen_server:start_link(?MODULE, {Queue, Concurrency}, []).

-doc false.
init({Queue, Concurrency}) ->
    process_flag(trap_exit, true),
    Pool = shigoto_config:pool(),
    Weights = shigoto_config:queue_weights(),
    Weight = maps:get(Queue, Weights, 1),
    FairQueues = shigoto_config:fair_queues(),
    Fair = lists:member(Queue, FairQueues),
    schedule_poll(),
    schedule_rescue(),
    {ok, #state{
        queue = Queue,
        concurrency = Concurrency,
        active = 0,
        pool = Pool,
        paused = false,
        shutting_down = false,
        executors = #{},
        weight = Weight,
        weight_counter = 0,
        fair = Fair
    }}.

-doc "Pause a queue — stops claiming new jobs but lets in-flight jobs finish.".
-spec pause(pid()) -> ok.
pause(Pid) ->
    gen_server:call(Pid, pause).

-doc "Resume a paused queue.".
-spec resume(pid()) -> ok.
resume(Pid) ->
    gen_server:call(Pid, resume).

-doc false.
handle_call(pause, _From, State) ->
    {reply, ok, State#state{paused = true}};
handle_call(resume, _From, State) ->
    {reply, ok, State#state{paused = false}};
handle_call(_Request, _From, State) ->
    {reply, ok, State}.

-doc false.
handle_cast({job_finished, _JobId, Pid}, #state{active = Active, executors = Execs} = State) ->
    NewExecs =
        case maps:find(Pid, Execs) of
            {ok, Ref} ->
                erlang:demonitor(Ref, [flush]),
                maps:remove(Pid, Execs);
            error ->
                Execs
        end,
    {noreply, State#state{active = Active - 1, executors = NewExecs}};
handle_cast(_Msg, State) ->
    {noreply, State}.

-doc false.
handle_info(poll, #state{shutting_down = true} = State) ->
    {noreply, State};
handle_info(poll, #state{paused = true} = State) ->
    schedule_poll(),
    {noreply, State};
handle_info(
    poll,
    #state{
        queue = Queue,
        concurrency = Conc,
        active = Active,
        pool = Pool,
        executors = Execs,
        weight = Weight,
        weight_counter = Counter,
        fair = Fair
    } = State
) ->
    NewCounter = Counter + 1,
    case should_poll(Weight, NewCounter) of
        false ->
            schedule_poll(),
            {noreply, State#state{weight_counter = NewCounter}};
        true ->
            Available = Conc - Active,
            {NewActive, NewExecs} =
                case Available > 0 of
                    true ->
                        case claim(Pool, Queue, Available, Fair) of
                            {ok, Jobs} ->
                                shigoto_telemetry:queue_poll(Queue, length(Jobs)),
                                lists:foreach(
                                    fun shigoto_telemetry:job_claimed/1,
                                    Jobs
                                ),
                                {Started, Execs1} = lists:foldl(
                                    fun(Job, {Count, AccExecs}) ->
                                        case
                                            shigoto_executor_sup:start_executor(
                                                Job, Pool, self()
                                            )
                                        of
                                            {ok, Pid} ->
                                                Ref = erlang:monitor(process, Pid),
                                                {Count + 1, AccExecs#{Pid => Ref}};
                                            {error, _} ->
                                                {Count, AccExecs}
                                        end
                                    end,
                                    {0, Execs},
                                    Jobs
                                ),
                                {Active + Started, Execs1};
                            {error, _} ->
                                {Active, Execs}
                        end;
                    false ->
                        {Active, Execs}
                end,
            schedule_poll(),
            {noreply, State#state{
                active = NewActive,
                executors = NewExecs,
                weight_counter = NewCounter
            }}
    end;
handle_info(rescue, #state{shutting_down = true} = State) ->
    {noreply, State};
handle_info(rescue, #state{pool = Pool} = State) ->
    StaleSeconds = stale_threshold(),
    _ = shigoto_repo:rescue_stale_jobs(Pool, StaleSeconds),
    schedule_rescue(),
    {noreply, State};
handle_info(
    {'DOWN', _Ref, process, Pid, _Reason}, #state{active = Active, executors = Execs} = State
) ->
    NewExecs = maps:remove(Pid, Execs),
    NewActive = max(0, Active - 1),
    {noreply, State#state{active = NewActive, executors = NewExecs}};
handle_info(_Info, State) ->
    {noreply, State}.

-doc false.
terminate(_Reason, #state{executors = Execs}) ->
    case map_size(Execs) of
        0 ->
            ok;
        _ ->
            Timeout = shigoto_config:shutdown_timeout(),
            wait_for_executors(Execs, Timeout)
    end.

%%----------------------------------------------------------------------
%% Internal
%%----------------------------------------------------------------------

claim(Pool, Queue, Limit, true) ->
    shigoto_repo:claim_jobs_fair(Pool, Queue, Limit);
claim(Pool, Queue, Limit, false) ->
    shigoto_repo:claim_jobs(Pool, Queue, Limit).

schedule_poll() ->
    erlang:send_after(shigoto_config:poll_interval(), self(), poll).

schedule_rescue() ->
    erlang:send_after(60000, self(), rescue).

stale_threshold() ->
    HeartbeatInterval = shigoto_config:heartbeat_interval(),
    HeartbeatInterval * 2 div 1000.

should_poll(1, _Counter) ->
    true;
should_poll(Weight, Counter) ->
    Counter rem Weight =:= 0.

wait_for_executors(Execs, _Timeout) when map_size(Execs) =:= 0 ->
    ok;
wait_for_executors(_Execs, Timeout) when Timeout =< 0 ->
    ok;
wait_for_executors(Execs, Timeout) ->
    T0 = erlang:monotonic_time(millisecond),
    receive
        {'DOWN', _Ref, process, Pid, _Reason} ->
            Elapsed = erlang:monotonic_time(millisecond) - T0,
            wait_for_executors(maps:remove(Pid, Execs), Timeout - Elapsed)
    after Timeout ->
        ok
    end.