mirror of
https://github.com/fluxerapp/fluxer
synced 2026-10-18 21:45:31 +09:00
125 lines
3.7 KiB
Erlang
125 lines
3.7 KiB
Erlang
%% SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
-module(session_manager_shards).
|
|
-typing([eqwalizer]).
|
|
|
|
-export([
|
|
ensure_shard/2,
|
|
restart_shard/2,
|
|
start_shards/1,
|
|
ensure_shard_table/0,
|
|
sync_shard_table/1,
|
|
delete_shard_table/0
|
|
]).
|
|
|
|
-export_type([session_id/0, shard/0, state/0]).
|
|
|
|
-define(SHARD_TABLE, session_manager_shard_table).
|
|
|
|
-type session_id() :: binary().
|
|
-type shard() :: #{pid := pid(), ref := reference()}.
|
|
-type state() :: #{shards := #{non_neg_integer() => shard()}, shard_count := pos_integer()}.
|
|
|
|
-spec ensure_shard(session_id(), state()) -> {non_neg_integer(), state()}.
|
|
ensure_shard(SessionId, State) ->
|
|
Count = maps:get(shard_count, State),
|
|
Shards = maps:get(shards, State),
|
|
Index = session_manager_routing:select_shard(SessionId, Count),
|
|
case maps:get(Index, Shards, undefined) of
|
|
undefined ->
|
|
{_Result, NewState} = restart_shard(Index, State),
|
|
{Index, NewState};
|
|
#{pid := Pid} ->
|
|
ensure_live_shard(SessionId, Index, Pid, State)
|
|
end.
|
|
|
|
-spec ensure_live_shard(session_id(), non_neg_integer(), pid(), state()) ->
|
|
{non_neg_integer(), state()}.
|
|
ensure_live_shard(_SessionId, Index, Pid, State) ->
|
|
case process_liveness:is_alive(Pid) of
|
|
true ->
|
|
{Index, State};
|
|
false ->
|
|
{_Result, NewState} = restart_shard(Index, State),
|
|
{Index, NewState}
|
|
end.
|
|
|
|
-spec restart_shard(non_neg_integer(), state()) -> {shard() | {error, term()}, state()}.
|
|
restart_shard(Index, State) ->
|
|
case start_shard(Index) of
|
|
{ok, Shard} ->
|
|
Shards = maps:get(shards, State),
|
|
NewState = State#{shards := Shards#{Index => Shard}},
|
|
sync_shard_table(NewState),
|
|
{Shard, NewState};
|
|
{error, Reason} ->
|
|
Shards = maps:get(shards, State),
|
|
NewState = State#{shards := maps:remove(Index, Shards)},
|
|
{{error, Reason}, NewState}
|
|
end.
|
|
|
|
-spec start_shards(pos_integer()) -> #{non_neg_integer() => shard()}.
|
|
start_shards(Count) ->
|
|
lists:foldl(
|
|
fun maybe_add_shard/2,
|
|
#{},
|
|
lists:seq(0, Count - 1)
|
|
).
|
|
|
|
-spec maybe_add_shard(non_neg_integer(), #{non_neg_integer() => shard()}) ->
|
|
#{non_neg_integer() => shard()}.
|
|
maybe_add_shard(Index, Acc) ->
|
|
case start_shard(Index) of
|
|
{ok, Shard} -> Acc#{Index => Shard};
|
|
{error, _Reason} -> Acc
|
|
end.
|
|
|
|
-spec ensure_shard_table() -> ok.
|
|
ensure_shard_table() ->
|
|
case ets:whereis(?SHARD_TABLE) of
|
|
undefined ->
|
|
_ = ets:new(?SHARD_TABLE, [named_table, public, set, {read_concurrency, true}]),
|
|
ok;
|
|
_ ->
|
|
ok
|
|
end.
|
|
|
|
-spec sync_shard_table(state()) -> ok.
|
|
sync_shard_table(State) ->
|
|
ensure_shard_table(),
|
|
_ = ets:delete_all_objects(?SHARD_TABLE),
|
|
ShardCount = maps:get(shard_count, State),
|
|
ets:insert(?SHARD_TABLE, {shard_count, ShardCount}),
|
|
Shards = maps:get(shards, State),
|
|
maps:foreach(
|
|
fun(Index, #{pid := Pid}) -> put_shard_pid(Index, Pid) end,
|
|
Shards
|
|
),
|
|
ok.
|
|
|
|
-spec delete_shard_table() -> ok.
|
|
delete_shard_table() ->
|
|
try ets:delete(?SHARD_TABLE) of
|
|
true -> ok
|
|
catch
|
|
error:badarg -> ok
|
|
end.
|
|
|
|
-spec start_shard(non_neg_integer()) -> {ok, shard()} | {error, term()}.
|
|
start_shard(Index) ->
|
|
case session_manager_shard:start_link(Index) of
|
|
{ok, Pid} ->
|
|
Ref = erlang:monitor(process, Pid),
|
|
put_shard_pid(Index, Pid),
|
|
{ok, #{pid => Pid, ref => Ref}};
|
|
ignore ->
|
|
{error, ignore};
|
|
Error ->
|
|
Error
|
|
end.
|
|
|
|
-spec put_shard_pid(non_neg_integer(), pid()) -> ok.
|
|
put_shard_pid(Index, Pid) ->
|
|
ets:insert(?SHARD_TABLE, {{shard_pid, Index}, Pid}),
|
|
ok.
|