fluxer/fluxer_gateway/src/guild/guild_counts_cache_shard_mgmt.erl

%% SPDX-License-Identifier: AGPL-3.0-or-later

-module(guild_counts_cache_shard_mgmt).
-typing([eqwalizer]).

-include_lib("fluxer_gateway/include/timeout_config.hrl").

-export([
    start_shards/2,
    start_shard/1,
    restart_shard/2,
    find_shard_by_ref/2,
    find_shard_by_pid/2,
    ensure_shard/2,
    ensure_shard_for_index/2,
    select_shard/2,
    forward_call/3,
    call_shard/3,
    shard_pid/2,
    determine_shard_count/0
]).

-export([
    rebalance_retry_timer_ref/1,
    set_rebalance_retry_timer/2,
    refresh_rebalance_retry_timer/1,
    ensure_rebalance_retry_timer/1,
    cancel_rebalance_retry_timer/1,
    ensure_state_fields/1
]).

-type shard() :: #{pid := pid(), ref := reference()}.

-export_type([shard/0]).

-define(PENDING_HANDOFF_RETRY_MS, 1000).

-spec determine_shard_count() -> {pos_integer(), atom() | auto}.
determine_shard_count() ->
    clustered_ets_cache:determine_shard_count([guild_counts_cache_shards, guild_shards]).

-spec start_shards(pos_integer(), #{}) -> #{non_neg_integer() => shard()}.
start_shards(Count, Acc) ->
    lists:foldl(
        fun start_shard_into/2,
        Acc,
        lists:seq(0, Count - 1)
    ).

-spec start_shard_into(non_neg_integer(), #{non_neg_integer() => shard()}) ->
    #{non_neg_integer() => shard()}.
start_shard_into(Index, MapAcc) ->
    case start_shard(Index) of
        {ok, Shard} -> MapAcc#{Index => Shard};
        {error, _Reason} -> MapAcc
    end.

-spec start_shard(non_neg_integer()) -> {ok, shard()} | {error, term()}.
start_shard(Index) ->
    case guild_counts_cache_shard:start_link(Index) of
        {ok, Pid} ->
            Ref = erlang:monitor(process, Pid),
            {ok, #{pid => Pid, ref => Ref}};
        ignore ->
            {error, ignore};
        Error ->
            Error
    end.

-spec restart_shard(non_neg_integer(), map()) -> {shard(), map()}.
restart_shard(Index, State) ->
    case start_shard(Index) of
        {ok, Shard} ->
            store_shard(Index, Shard, State);
        {error, _Reason} ->
            Dummy = make_dummy_shard(),
            store_shard(Index, Dummy, State)
    end.

-spec store_shard(non_neg_integer(), shard(), map()) -> {shard(), map()}.
store_shard(Index, Shard, State) ->
    Shards = maps:get(shards, State),
    {Shard, State#{shards := Shards#{Index => Shard}}}.

-spec make_dummy_shard() -> shard().
make_dummy_shard() ->
    DummyPid = spawn(fun() -> ok = gateway_retry_timer:wait(50) end),
    #{pid => DummyPid, ref => make_ref()}.

-spec forward_call(term(), term(), map()) -> {term(), map()}.
forward_call(Key, Request, State) ->
    {Index, State1} = ensure_shard(Key, State),
    call_shard(Index, Request, State1).

-spec call_shard(non_neg_integer(), term(), map()) -> {term(), map()}.
call_shard(Index, Request, State) ->
    Pid = shard_pid(Index, State),
    try gen_server:call(Pid, Request, ?DEFAULT_GEN_SERVER_TIMEOUT) of
        Reply -> {Reply, State}
    catch
        exit:_ ->
            {_Shard, State1} = restart_shard(Index, State),
            call_shard_after_restart(Index, Request, State1)
    end.

-spec call_shard_after_restart(non_neg_integer(), term(), map()) -> {term(), map()}.
call_shard_after_restart(Index, Request, State) ->
    Pid = shard_pid(Index, State),
    try gen_server:call(Pid, Request, ?DEFAULT_GEN_SERVER_TIMEOUT) of
        Reply -> {Reply, State}
    catch
        exit:_ -> {error, State}
    end.

-spec shard_pid(non_neg_integer(), map()) -> pid().
shard_pid(Index, State) ->
    Shards = maps:get(shards, State),
    Shard = maps:get(Index, Shards),
    maps:get(pid, Shard).

-spec ensure_shard(term(), map()) -> {non_neg_integer(), map()}.
ensure_shard(Key, State) ->
    Count = maps:get(shard_count, State),
    Index = select_shard(Key, Count),
    ensure_shard_for_index(Index, State).

-spec ensure_shard_for_index(non_neg_integer(), map()) -> {non_neg_integer(), map()}.
ensure_shard_for_index(Index, State) ->
    Shards = maps:get(shards, State),
    case maps:get(Index, Shards, undefined) of
        undefined ->
            {_Shard, NewState} = restart_shard(Index, State),
            {Index, NewState};
        #{pid := Pid} when is_pid(Pid) ->
            handle_shard_liveness(Index, Pid, State)
    end.

-spec handle_shard_liveness(non_neg_integer(), pid(), map()) ->
    {non_neg_integer(), map()}.
handle_shard_liveness(Index, Pid, State) ->
    case process_liveness:is_alive(Pid) of
        true ->
            {Index, State};
        false ->
            {_Shard, NewState} = restart_shard(Index, State),
            {Index, NewState}
    end.

-spec select_shard(term(), pos_integer()) -> non_neg_integer().
select_shard(Key, Count) when Count > 0 ->
    clustered_ets_cache:select_shard(Key, Count).

-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard()}) ->
    {ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
    shard_utils:find_shard_by_ref(Ref, Shards).

-spec find_shard_by_pid(pid(), #{non_neg_integer() => shard()}) ->
    {ok, non_neg_integer()} | not_found.
find_shard_by_pid(Pid, Shards) ->
    shard_utils:find_shard_by_pid(Pid, Shards).

-spec rebalance_retry_timer_ref(map()) -> reference() | undefined.
rebalance_retry_timer_ref(State) ->
    case maps:get(rebalance_retry_timer, State, undefined) of
        TimerRef when is_reference(TimerRef) -> TimerRef;
        _ -> undefined
    end.

-spec set_rebalance_retry_timer(reference() | undefined, map()) -> map().
set_rebalance_retry_timer(TimerRef, State) ->
    State#{rebalance_retry_timer => TimerRef}.

-spec ensure_rebalance_retry_timer(map()) -> map().
ensure_rebalance_retry_timer(State) ->
    case rebalance_retry_timer_ref(State) of
        TimerRef when is_reference(TimerRef) -> State;
        undefined ->
            NewRef = erlang:start_timer(?PENDING_HANDOFF_RETRY_MS, self(), rebalance_retry),
            set_rebalance_retry_timer(NewRef, State)
    end.

-spec cancel_rebalance_retry_timer(map()) -> map().
cancel_rebalance_retry_timer(State) ->
    case rebalance_retry_timer_ref(State) of
        TimerRef when is_reference(TimerRef) ->
            _ = erlang:cancel_timer(TimerRef),
            set_rebalance_retry_timer(undefined, State);
        undefined ->
            set_rebalance_retry_timer(undefined, State)
    end.

-spec refresh_rebalance_retry_timer(map()) -> map().
refresh_rebalance_retry_timer(State) ->
    case guild_counts_cache_remote:pending_handoff_count_from_state(State) of
        0 -> cancel_rebalance_retry_timer(State);
        _Count -> ensure_rebalance_retry_timer(State)
    end.

-spec ensure_state_fields(map()) -> map().
ensure_state_fields(State) when is_map(State) ->
    Pending = guild_counts_cache_remote:pending_handoffs(State),
    TimerRef = rebalance_retry_timer_ref(State),
    State1 = guild_counts_cache_remote:set_pending_handoffs(Pending, State),
    set_rebalance_retry_timer(TimerRef, State1).