refactor(gateway): share shard helpers and test fixtures (#3381)

This commit is contained in:
Hampus authored and GitHub committed 2026-10-10 23:04:30 +02:00
1 parent f22261f218
commit e7b13fd19f
25 files changed
+101 -2021

No files matched your search

@@ -153,26 +153,12 @@ select_shard(Key, Count) when Count > 0 ->
-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
maps:fold(
fun
(Index, #{ref := R}, _) when R =:= Ref -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
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) ->
maps:fold(
fun
(Index, #{pid := P}, _) when P =:= Pid -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
shard_utils:find_shard_by_pid(Pid, Shards).
-spec rebalance_retry_timer_ref(map()) -> reference() | undefined.
rebalance_retry_timer_ref(State) ->
@@ -25,19 +25,11 @@
-spec determine_shard_count() -> {pos_integer(), configured | auto}.
determine_shard_count() ->
case fluxer_gateway_env:get(guild_shards) of
Value when is_integer(Value), Value > 0 ->
{Value, configured};
_ ->
{default_shard_count(), auto}
end.
shard_utils:determine_shard_count(guild_shards).
-spec default_shard_count() -> pos_integer().
default_shard_count() ->
shard_utils:max_positive([
erlang:system_info(logical_processors_available),
erlang:system_info(schedulers_online)
]).
clustered_ets_cache:default_shard_count().
-spec start_shards(pos_integer()) -> #{non_neg_integer() => shard_map()}.
start_shards(Count) ->
@@ -109,12 +101,12 @@ normalize_grouped_ids(Groups) ->
-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard_map()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
find_shard_by(fun(#{ref := R}) -> R =:= Ref end, Shards).
shard_utils:find_shard_by_ref(Ref, Shards).
-spec find_shard_by_pid(pid(), #{non_neg_integer() => shard_map()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_pid(Pid, Shards) ->
find_shard_by(fun(#{pid := P}) -> P =:= Pid end, Shards).
shard_utils:find_shard_by_pid(Pid, Shards).
-spec start_shard_into_map(non_neg_integer(), #{non_neg_integer() => shard_map()}) ->
#{non_neg_integer() => shard_map()}.
@@ -132,26 +124,3 @@ ensure_live_shard(Index, Pid, State) ->
false ->
restart_shard_or_fail(Index, State)
end.
-spec find_shard_by(fun((shard_map()) -> boolean()), #{non_neg_integer() => shard_map()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by(Pred, Shards) ->
maps:fold(
fun(Index, ShardMap, Acc) -> find_matching_shard(Pred, Index, ShardMap, Acc) end,
not_found,
Shards
).
-spec find_matching_shard(
fun((shard_map()) -> boolean()),
non_neg_integer(),
shard_map(),
{ok, non_neg_integer()} | not_found
) -> {ok, non_neg_integer()} | not_found.
find_matching_shard(_Pred, _Index, _ShardMap, {ok, _Found} = Acc) ->
Acc;
find_matching_shard(Pred, Index, ShardMap, not_found) ->
case Pred(ShardMap) of
true -> {ok, Index};
false -> not_found
end.
@@ -323,7 +323,7 @@ maybe_add_last_message_id(RawId, RawMsgId, UserId, Member, State, Acc) ->
binary(), binary(), term(), integer(), map(), guild_state(), map()
) -> map().
maybe_add_last_message_for_channel(IdBin, MsgIdBin, RawId, UserId, Member, State, Acc) ->
case parse_snowflake(<<"id">>, RawId) of
case guild_dispatch_decorate:parse_snowflake(<<"id">>, RawId) of
ChId when is_integer(ChId) ->
maybe_add_visible_last_message(IdBin, MsgIdBin, UserId, ChId, Member, State, Acc);
undefined ->
@@ -345,13 +345,6 @@ maybe_add_visible_last_message(IdBin, MsgId, UserId, ChId, Member, State, Acc) -
compute_voice_state_updates(Current, Previous, GuildId) ->
guild_passive_sync_voice:compute_voice_state_updates(Current, Previous, GuildId).
-spec parse_snowflake(binary(), term()) -> integer() | undefined.
parse_snowflake(FieldName, Value) ->
case validation:validate_snowflake(FieldName, Value) of
{ok, Id} -> Id;
{error, _, _} -> undefined
end.
-spec snowflake_binary(binary(), term()) -> binary() | undefined.
snowflake_binary(FieldName, Value) ->
case validation:validate_snowflake(FieldName, Value) of
@@ -150,7 +150,7 @@ moderator_rules_test() ->
).
shared_case_table_test() ->
{ok, Contents} = file:read_file("test/thread_permission_cases.json"),
{ok, Contents} = file:read_file("../packages/constants/src/ThreadPermissionCases.json"),
Cases = [
Case
|| #{<<"fn">> := Fn} = Case <- json:decode(Contents),
@@ -39,7 +39,7 @@ ensure_tables() ->
-spec upsert_voice_state(voice_state()) -> ok.
upsert_voice_state(VoiceState) when is_map(VoiceState) ->
ensure_tables(),
ConnectionId = normalize_optional_binary(
ConnectionId = guild_voice_connection_normalize:normalize_optional_binary(
get_optional_field(VoiceState, <<"connection_id">>, connection_id)
),
upsert_by_connection(ConnectionId, VoiceState);
@@ -51,17 +51,17 @@ upsert_by_connection(undefined, _VoiceState) ->
ok;
upsert_by_connection(ConnectionId, VoiceState) ->
case
normalize_optional_binary(
guild_voice_connection_normalize:normalize_optional_binary(
get_optional_field(VoiceState, <<"channel_id">>, channel_id, null)
)
of
undefined ->
remove_connection(ConnectionId);
_ChannelId ->
RegionId = normalize_optional_binary(
RegionId = guild_voice_connection_normalize:normalize_optional_binary(
get_optional_field(VoiceState, <<"region_id">>, region_id)
),
ServerId = normalize_optional_binary(
ServerId = guild_voice_connection_normalize:normalize_optional_binary(
get_optional_field(VoiceState, <<"server_id">>, server_id)
),
upsert_connection_metadata(ConnectionId, RegionId, ServerId)
@@ -70,7 +70,7 @@ upsert_by_connection(ConnectionId, VoiceState) ->
-spec remove_connection(binary() | term()) -> ok.
remove_connection(ConnectionId) ->
ensure_tables(),
case normalize_optional_binary(ConnectionId) of
case guild_voice_connection_normalize:normalize_optional_binary(ConnectionId) of
undefined ->
ok;
NormalizedConnectionId ->
@@ -214,19 +214,6 @@ compare_count_rows({LeftId, LeftCount}, {RightId, RightCount}) ->
false -> LeftCount > RightCount
end.
-spec normalize_optional_binary(term()) -> binary() | undefined.
normalize_optional_binary(undefined) ->
undefined;
normalize_optional_binary(null) ->
undefined;
normalize_optional_binary(Value) when is_binary(Value), byte_size(Value) > 0 -> Value;
normalize_optional_binary(Value) when is_binary(Value) -> undefined;
normalize_optional_binary(Value) when is_integer(Value) -> integer_to_binary(Value);
normalize_optional_binary(Value) when is_list(Value) ->
guild_voice_connection_normalize:normalize_optional_binary(Value);
normalize_optional_binary(_) ->
undefined.
-spec get_optional_field(map(), term(), term()) -> term().
get_optional_field(Map, PrimaryKey, FallbackKey) ->
get_optional_field(Map, PrimaryKey, FallbackKey, undefined).
@@ -159,7 +159,7 @@ maybe_put_live_voice_state(FallbackConnId, DefaultChanId, VoiceState, Acc) ->
VoiceState,
maps:get(connection_id, VoiceState, FallbackConnId)
),
ConnectionId = normalize_optional_binary(RawConnId),
ConnectionId = guild_voice_connection_normalize:normalize_optional_binary(RawConnId),
case ConnectionId of
undefined ->
Acc;
@@ -175,7 +175,7 @@ maybe_put_default_channel_id(VoiceState, DefaultChannelId) ->
VoiceState,
maps:get(channel_id, VoiceState, undefined)
),
case normalize_optional_binary(ExistingId) of
case guild_voice_connection_normalize:normalize_optional_binary(ExistingId) of
undefined ->
apply_default_channel_id(VoiceState, DefaultChannelId);
ChannelId ->
@@ -184,27 +184,11 @@ maybe_put_default_channel_id(VoiceState, DefaultChannelId) ->
-spec apply_default_channel_id(voice_state(), term()) -> voice_state().
apply_default_channel_id(VoiceState, DefaultChannelId) ->
case normalize_optional_binary(DefaultChannelId) of
case guild_voice_connection_normalize:normalize_optional_binary(DefaultChannelId) of
undefined -> VoiceState;
ChannelId -> VoiceState#{<<"channel_id">> => ChannelId}
end.
-spec normalize_optional_binary(term()) -> binary() | undefined.
normalize_optional_binary(undefined) ->
undefined;
normalize_optional_binary(null) ->
undefined;
normalize_optional_binary(Value) when is_binary(Value), byte_size(Value) > 0 ->
Value;
normalize_optional_binary(Value) when is_binary(Value) ->
undefined;
normalize_optional_binary(Value) when is_integer(Value) ->
integer_to_binary(Value);
normalize_optional_binary(Value) when is_list(Value) ->
guild_voice_connection_normalize:normalize_optional_binary(Value);
normalize_optional_binary(_) ->
undefined.
-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
+3 -33
View File
@@ -44,7 +44,7 @@ diagnostic_info() ->
init([]) ->
process_flag(trap_exit, true),
erlang:process_flag(fullsweep_after, 10),
{ShardCount, _Source} = determine_shard_count(presence_bus_shards),
{ShardCount, _Source} = shard_utils:determine_shard_count(presence_bus_shards),
Shards = start_shards(ShardCount, #{}),
{ok, #{shards => Shards, shard_count => ShardCount}, hibernate}.
@@ -127,22 +127,6 @@ code_change(_OldVsn, State, _Extra) ->
erlang:garbage_collect(),
{ok, State}.
-spec determine_shard_count(atom()) -> {pos_integer(), configured | auto}.
determine_shard_count(ConfigKey) ->
case fluxer_gateway_env:get(ConfigKey) of
Value when is_integer(Value), Value > 0 ->
{Value, configured};
_ ->
{default_shard_count(), auto}
end.
-spec default_shard_count() -> pos_integer().
default_shard_count() ->
shard_utils:max_positive([
erlang:system_info(logical_processors_available),
erlang:system_info(schedulers_online)
]).
-spec start_shards(pos_integer(), #{}) -> #{non_neg_integer() => shard()}.
start_shards(Count, Acc) ->
lists:foldl(fun start_shard_acc/2, Acc, lists:seq(0, Count - 1)).
@@ -238,26 +222,12 @@ select_shard(Key, Count) when Count > 0 ->
-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
maps:fold(
fun
(Index, #{ref := R}, _) when R =:= Ref -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
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) ->
maps:fold(
fun
(Index, #{pid := P}, _) when P =:= Pid -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
shard_utils:find_shard_by_pid(Pid, Shards).
-spec safe_gen_server_call(pid() | atom(), term(), timeout()) -> term().
safe_gen_server_call(Server, Request, Timeout) ->
@@ -193,7 +193,7 @@ add_user_to_owner_group(UserId, Acc) ->
-spec local_shard_count() -> pos_integer().
local_shard_count() ->
{Count, _Source} = determine_shard_count(presence_cache_shards),
{Count, _Source} = presence_cache_shards:determine_count(presence_cache_shards),
Count.
-spec select_shard(term(), pos_integer()) -> non_neg_integer().
@@ -271,14 +271,6 @@ accumulate_ets_lookup(TableName, UserId, AccMap) ->
error:badarg -> AccMap
end.
-spec determine_shard_count(atom()) -> {pos_integer(), configured | auto}.
determine_shard_count(ConfigKey) ->
case clustered_ets_cache:determine_shard_count([ConfigKey]) of
{Count, ConfigKey} -> {Count, configured};
{Count, auto} -> {Count, auto};
{Count, _Source} -> {Count, configured}
end.
-spec normalize_shard_groups([{non_neg_integer(), [term()]}]) ->
[{non_neg_integer(), [integer()]}].
normalize_shard_groups(Groups) ->
@@ -66,7 +66,7 @@ handle_down_by_ref(Ref, State) ->
-spec handle_down_by_pid(pid(), state()) -> state().
handle_down_by_pid(Pid, State) ->
Shards = maps:get(shards, State),
case find_by_pid(Pid, Shards) of
case shard_utils:find_shard_by_pid(Pid, Shards) of
{ok, Index} -> restart_state(Index, State);
not_found -> State
end.
@@ -196,14 +196,7 @@ memory_stats(State) ->
-spec find_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_by_ref(Ref, Shards) ->
maps:fold(
fun
(Index, #{ref := R}, _) when R =:= Ref -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
shard_utils:find_shard_by_ref(Ref, Shards).
-spec start(non_neg_integer()) -> {ok, shard()} | {error, term()}.
start(Index) ->
@@ -306,18 +299,6 @@ merge_snapshot_reply(Snapshot, AccMap) when is_map(Snapshot) ->
merge_snapshot_reply(_, AccMap) ->
AccMap.
-spec find_by_pid(pid(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_by_pid(Pid, Shards) ->
maps:fold(
fun
(Index, #{pid := P}, _) when P =:= Pid -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
-spec error_reply(term()) -> term().
error_reply({put, _, _}) -> {error, unavailable};
error_reply({delete, _}) -> {error, unavailable};
@@ -36,10 +36,7 @@
-spec determine_count() -> {pos_integer(), configured | auto}.
determine_count() ->
case fluxer_gateway_env:get(presence_shards) of
Value when is_integer(Value), Value > 0 -> {Value, configured};
_ -> {default_count(), auto}
end.
shard_utils:determine_shard_count(presence_shards).
-spec start(non_neg_integer()) -> {ok, shard()} | {error, term()}.
start(Index) ->
@@ -157,33 +154,16 @@ pid_from_table(Key) ->
-spec find_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_by_ref(Ref, Shards) ->
maps:fold(
fun
(Index, #{ref := R}, _) when R =:= Ref -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
shard_utils:find_shard_by_ref(Ref, Shards).
-spec find_by_pid(pid(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_by_pid(Pid, Shards) ->
maps:fold(
fun
(Index, #{pid := P}, _) when P =:= Pid -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
shard_utils:find_shard_by_pid(Pid, Shards).
-spec default_count() -> pos_integer().
default_count() ->
shard_utils:max_positive([
erlang:system_info(logical_processors_available),
erlang:system_info(schedulers_online)
]).
clustered_ets_cache:default_shard_count().
-spec ensure_shard(user_id(), state()) -> {non_neg_integer(), state()}.
ensure_shard(Key, State) ->
@@ -269,7 +269,7 @@ push_eligible(State) ->
-spec build_push_create_params(user_id(), map()) -> map() | undefined.
build_push_create_params(UserId, Data) ->
AuthorIdBin = maps:get(<<"id">>, maps:get(<<"author">>, Data, #{}), undefined),
case parse_snowflake(<<"author_id">>, AuthorIdBin) of
case guild_dispatch_decorate:parse_snowflake(<<"author_id">>, AuthorIdBin) of
undefined ->
undefined;
AuthorId ->
@@ -439,14 +439,7 @@ session_holds_push(Session) ->
-spec extract_snowflake(binary(), map()) -> integer() | undefined.
extract_snowflake(FieldName, Data) ->
parse_snowflake(FieldName, maps:get(FieldName, Data, undefined)).
-spec parse_snowflake(binary(), term()) -> integer() | undefined.
parse_snowflake(FieldName, Value) ->
case validation:validate_snowflake(FieldName, Value) of
{ok, Id} -> Id;
{error, _, _} -> undefined
end.
guild_dispatch_decorate:parse_snowflake(FieldName, maps:get(FieldName, Data, undefined)).
-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
+1 -17
View File
@@ -96,7 +96,7 @@ init([]) ->
erlang:process_flag(fullsweep_after, 0),
_ = fluxer_gateway_env:load(),
session_manager_shards:ensure_shard_table(),
{ShardCount, _Source} = determine_shard_count(),
{ShardCount, _Source} = shard_utils:determine_shard_count(session_shards),
Shards = session_manager_shards:start_shards(ShardCount),
State = #{shards => Shards, shard_count => ShardCount},
session_manager_shards:sync_shard_table(State),
@@ -316,19 +316,3 @@ atom_list([Value | Rest], Acc) when is_atom(Value) ->
atom_list(Rest, [Value | Acc]);
atom_list(_, _) ->
error.
-spec determine_shard_count() -> {pos_integer(), configured | auto}.
determine_shard_count() ->
case fluxer_gateway_env:get(session_shards) of
Value when is_integer(Value), Value > 0 ->
{Value, configured};
_ ->
{default_shard_count(), auto}
end.
-spec default_shard_count() -> pos_integer().
default_shard_count() ->
shard_utils:max_positive([
erlang:system_info(logical_processors_available),
erlang:system_info(schedulers_online)
]).
@@ -93,28 +93,12 @@ safe_shard_call(Pid, Request) ->
-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
maps:fold(
fun
(_Index, _Shard, {ok, _} = Found) -> Found;
(Index, #{ref := ExistingRef}, not_found) when ExistingRef =:= Ref -> {ok, Index};
(_Index, _Shard, not_found) -> not_found
end,
not_found,
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) ->
maps:fold(
fun
(_Index, _Shard, {ok, _} = Found) -> Found;
(Index, #{pid := ExistingPid}, not_found) when ExistingPid =:= Pid -> {ok, Index};
(_Index, _Shard, not_found) -> not_found
end,
not_found,
Shards
).
shard_utils:find_shard_by_pid(Pid, Shards).
-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
@@ -188,6 +172,21 @@ aggregate_handoff_to_topology_marks_invalid_reply_as_failed_test() ->
aggregate_handoff_to_topology(['peer@host'], State),
InvalidShardPid ! stop.
aggregate_handoff_to_topology_counts_unavailable_and_bare_replies_test() ->
ExitedShardPid = spawn(fun() -> ok end),
ok = gateway_retry_timer:wait(5),
BareShardPid = spawn(fun() -> shard_bare_reply_test_loop(2) end),
State = #{
shards => #{
0 => #{pid => ExitedShardPid, ref => make_ref()},
1 => #{pid => BareShardPid, ref => make_ref()}
},
shard_count => 2
},
{#{attempted := 3, handed_off := 2}, _State} =
aggregate_handoff_to_topology(['peer@host'], State),
BareShardPid ! stop.
shard_reconnect_test_loop(DrainCount) ->
receive
{'$gen_call', From, reconnect_drain} ->
@@ -217,10 +216,13 @@ shard_handoff_test_loop(Result) ->
end.
shard_invalid_reply_test_loop() ->
shard_bare_reply_test_loop(invalid_reply).
shard_bare_reply_test_loop(Reply) ->
receive
{'$gen_call', From, _Request} ->
gen_server:reply(From, invalid_reply),
shard_invalid_reply_test_loop();
gen_server:reply(From, Reply),
shard_bare_reply_test_loop(Reply);
stop ->
ok
after 5000 ->
+40
View File
@@ -5,6 +5,9 @@
-export([
max_positive/1,
determine_shard_count/1,
find_shard_by_ref/2,
find_shard_by_pid/2,
safe_gen_call/3,
safe_gen_call_detailed/3,
safe_gen_call_remote/3,
@@ -14,10 +17,47 @@
check_mailbox_pressure/1
]).
-export_type([shard/0]).
-type shard() :: #{pid := pid(), ref := reference()}.
-spec max_positive(list()) -> pos_integer().
max_positive(Candidates) ->
lists:max([C || C <- Candidates, is_integer(C), C > 0] ++ [1]).
-spec determine_shard_count(atom()) -> {pos_integer(), configured | auto}.
determine_shard_count(ConfigKey) ->
case fluxer_gateway_env:get(ConfigKey) of
Value when is_integer(Value), Value > 0 ->
{Value, configured};
_ ->
{clustered_ets_cache:default_shard_count(), auto}
end.
-spec find_shard_by_ref(reference(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_ref(Ref, Shards) ->
maps:fold(
fun
(Index, #{ref := R}, _) when R =:= Ref -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
-spec find_shard_by_pid(pid(), #{non_neg_integer() => shard()}) ->
{ok, non_neg_integer()} | not_found.
find_shard_by_pid(Pid, Shards) ->
maps:fold(
fun
(Index, #{pid := P}, _) when P =:= Pid -> {ok, Index};
(_, _, Acc) -> Acc
end,
not_found,
Shards
).
-spec safe_gen_call(pid() | atom(), term(), pos_integer()) -> term().
safe_gen_call(Pid, Request, Timeout) ->
try
@@ -1,98 +0,0 @@
%% SPDX-License-Identifier: AGPL-3.0-or-later
-module(gateway_dispatch_relay_ordering_tests).
-typing([eqwalizer]).
-include_lib("eunit/include/eunit.hrl").
-define(STATE_KEY, {gateway_dispatch_relay, state}).
-define(EVENT_COUNT, 8).
dispatch_preserves_order_when_worker_is_backpressured_test_() ->
{timeout, 30, fun dispatch_preserves_order_when_worker_is_backpressured/0}.
dispatch_many_preserves_order_when_worker_is_backpressured_test_() ->
{timeout, 30, fun dispatch_many_preserves_order_when_worker_is_backpressured/0}.
dispatch_preserves_order_when_worker_is_backpressured() ->
assert_ordered_delivery(fun(SessionPid, N) ->
gateway_dispatch_relay:dispatch(SessionPid, relay_order_event, #{<<"n">> => N})
end).
dispatch_many_preserves_order_when_worker_is_backpressured() ->
assert_ordered_delivery(fun(SessionPid, N) ->
gateway_dispatch_relay:dispatch_many([SessionPid], relay_order_event, #{<<"n">> => N})
end).
assert_ordered_delivery(SendFun) ->
Worker = gateway_dispatch_relay_batch:start_worker(0),
try
with_relay_workers([Worker], fun() -> ordered_delivery(Worker, SendFun) end)
after
stop_worker(Worker)
end.
ordered_delivery(Worker, SendFun) ->
Ref = make_ref(),
Parent = self(),
SessionPid = spawn_link(fun() -> session_loop(Parent, Ref) end),
MaxQueue = gateway_dispatch_relay_batch:max_queue(),
ok = sys:suspend(Worker),
fill_queue(Worker, MaxQueue - ?EVENT_COUNT),
lists:foreach(fun(N) -> ok = SendFun(SessionPid, N) end, lists:seq(1, ?EVENT_COUNT)),
?assert(gateway_dispatch_relay_batch:message_queue_len(Worker) >= MaxQueue),
ok = sys:resume(Worker),
Observed = collect_observed(Ref, ?EVENT_COUNT, []),
SessionPid ! stop,
?assertEqual(lists:seq(1, ?EVENT_COUNT), Observed).
fill_queue(_Worker, 0) ->
ok;
fill_queue(Worker, Remaining) ->
Worker ! relay_order_filler,
fill_queue(Worker, Remaining - 1).
session_loop(Parent, Ref) ->
receive
{'$gen_cast', {dispatch, relay_order_event, #{<<"n">> := N}}} ->
Parent ! {relay_order_received, Ref, N},
session_loop(Parent, Ref);
stop ->
ok
after 30000 ->
ok
end.
collect_observed(_Ref, 0, Acc) ->
lists:reverse(Acc);
collect_observed(Ref, Remaining, Acc) ->
receive
{relay_order_received, Ref, N} -> collect_observed(Ref, Remaining - 1, [N | Acc])
after 10000 ->
?assert(false, {relay_order_timeout, Remaining, lists:reverse(Acc)})
end.
with_relay_workers(Workers, Fun) ->
Previous = persistent_term:get(?STATE_KEY, undefined),
persistent_term:put(?STATE_KEY, #{
workers => list_to_tuple(Workers),
shard_count => length(Workers)
}),
try
Fun()
after
restore_relay_state(Previous)
end.
restore_relay_state(undefined) ->
_ = persistent_term:erase(?STATE_KEY),
ok;
restore_relay_state(Previous) ->
persistent_term:put(?STATE_KEY, Previous).
stop_worker(Pid) ->
try gen_server:stop(Pid, normal, 5000) of
_ -> ok
catch
exit:_ -> ok
end.
@@ -1,131 +0,0 @@
%% SPDX-License-Identifier: AGPL-3.0-or-later
-module(session_manager_transfer_stress_tests).
-typing([eqwalizer]).
-include_lib("eunit/include/eunit.hrl").
aggregate_counts_many_shards_mixed_replies_test_() ->
{timeout, 15, fun aggregate_counts_many_shards_mixed_replies/0}.
aggregate_handoff_many_shards_mixed_replies_test_() ->
{timeout, 15, fun aggregate_handoff_many_shards_mixed_replies/0}.
aggregate_counts_many_shards_mixed_replies() ->
{Shards, LivePids, ExpectedCount} = build_count_shards(256),
State = #{shards => Shards, shard_count => 256},
try
{ExpectedCount, State} = session_manager_transfer:aggregate_counts(
reconnect_drain, State
)
after
stop_live_shards(LivePids)
end.
aggregate_handoff_many_shards_mixed_replies() ->
{Shards, LivePids, ExpectedResult} = build_handoff_shards(128),
State = #{shards => Shards, shard_count => 128},
try
{ExpectedResult, State} =
session_manager_transfer:aggregate_handoff_to_topology(['peer@host'], State)
after
stop_live_shards(LivePids)
end.
build_count_shards(Count) ->
lists:foldl(
fun(Index, {ShardAcc, PidAcc, ExpectedAcc}) ->
{Pid, MaybeLive, Expected} = count_shard(Index),
{
ShardAcc#{Index => #{pid => Pid, ref => make_ref()}},
add_live_pid(MaybeLive, Pid, PidAcc),
ExpectedAcc + Expected
}
end,
{#{}, [], 0},
lists:seq(1, Count)
).
count_shard(Index) ->
case Index rem 10 of
0 ->
dead_shard(0);
1 ->
{spawn_shard(invalid_reply), live, 0};
_ ->
{spawn_shard({ok, 1}), live, 1}
end.
build_handoff_shards(Count) ->
lists:foldl(
fun(Index, {ShardAcc, PidAcc, ExpectedAcc}) ->
{Pid, MaybeLive, Expected} = handoff_shard(Index),
{
ShardAcc#{Index => #{pid => Pid, ref => make_ref()}},
add_live_pid(MaybeLive, Pid, PidAcc),
sum_handoff(ExpectedAcc, Expected)
}
end,
{#{}, [], #{attempted => 0, handed_off => 0}},
lists:seq(1, Count)
).
handoff_shard(Index) ->
case Index rem 8 of
0 ->
dead_shard(#{attempted => 1, handed_off => 0});
1 ->
{spawn_shard(invalid_reply), live, #{attempted => 1, handed_off => 0}};
2 ->
{spawn_shard(2), live, #{attempted => 2, handed_off => 2}};
_ ->
{spawn_shard({ok, #{attempted => 3, handed_off => 2}}), live, #{
attempted => 3,
handed_off => 2
}}
end.
dead_shard(Expected) ->
Pid = spawn(fun() -> ok end),
wait_for_dead(Pid, 100),
{Pid, dead, Expected}.
spawn_shard(Reply) ->
spawn(fun() -> shard_loop(Reply) end).
shard_loop(Reply) ->
receive
{'$gen_call', From, _Request} ->
gen_server:reply(From, Reply),
shard_loop(Reply);
stop ->
ok
after 10000 ->
ok
end.
add_live_pid(live, Pid, Pids) ->
[Pid | Pids];
add_live_pid(dead, _Pid, Pids) ->
Pids.
stop_live_shards(Pids) ->
lists:foreach(fun(Pid) -> Pid ! stop end, Pids).
sum_handoff(Acc, Result) ->
#{
attempted => maps:get(attempted, Acc) + maps:get(attempted, Result),
handed_off => maps:get(handed_off, Acc) + maps:get(handed_off, Result)
}.
wait_for_dead(Pid, Attempts) ->
case is_process_alive(Pid) of
false ->
ok;
true when Attempts > 0 ->
timer:sleep(1),
wait_for_dead(Pid, Attempts - 1);
true ->
exit(Pid, kill),
ok
end.
File diff suppressed because it is too large. Load diff
-60
View File
@@ -1,60 +0,0 @@
// SPDX-License-Identifier: AGPL-3.0-or-later
pub use crate::constants::{clamp_size, parse_image_size};
#[cfg(test)]
mod tests {
use super::*;
use crate::constants::{AssetKind, DEFAULT_IMAGE_SIZE, IMAGE_SIZES};
#[test]
fn unparsable_sizes_fall_back_to_the_default() {
assert_eq!(DEFAULT_IMAGE_SIZE, parse_image_size(None));
assert_eq!(DEFAULT_IMAGE_SIZE, parse_image_size(Some("")));
assert_eq!(DEFAULT_IMAGE_SIZE, parse_image_size(Some("-1")));
assert_eq!(DEFAULT_IMAGE_SIZE, parse_image_size(Some("99999999999")));
assert_eq!(DEFAULT_IMAGE_SIZE, parse_image_size(Some("not-a-number")));
}
#[test]
fn parsable_sizes_snap_up_to_the_next_rung() {
assert_eq!(16, parse_image_size(Some("0")));
assert_eq!(640, parse_image_size(Some("640")));
assert_eq!(1024, parse_image_size(Some("641")));
assert_eq!(1024, parse_image_size(Some("777")));
assert_eq!(1024, parse_image_size(Some("1000")));
assert_eq!(16384, parse_image_size(Some("99999")));
for size in IMAGE_SIZES {
assert_eq!(*size, parse_image_size(Some(&size.to_string())));
}
assert_eq!(16384, parse_image_size(Some("16384")));
}
#[test]
fn an_off_ladder_size_never_serves_fewer_pixels_than_requested() {
for requested in 1..=4096u32 {
let served = parse_image_size(Some(&requested.to_string()));
assert!(
served >= requested,
"size={requested} served {served}, a silent downscale"
);
}
}
#[test]
fn snapping_is_idempotent() {
for raw in [0u32, 1, 17, 641, 777, 1000, 4097, 99999] {
let once = parse_image_size(Some(&raw.to_string()));
let twice = parse_image_size(Some(&once.to_string()));
assert_eq!(once, twice, "size={raw} did not settle");
}
}
#[test]
fn clamping_respects_kind_bounds() {
assert_eq!(128, clamp_size(1, AssetKind::Avatar));
assert_eq!(1024, clamp_size(4096, AssetKind::Avatar));
assert_eq!(480, clamp_size(16, AssetKind::Banner));
assert_eq!(512, clamp_size(4096, AssetKind::Emoji));
}
}
+3
View File
@@ -157,6 +157,9 @@ mod tests {
assert_eq!(16, parse_image_size(Some("0")));
assert_eq!(16384, parse_image_size(Some("99999")));
assert_eq!(128, parse_image_size(Some("not-a-number")));
assert_eq!(128, parse_image_size(Some("")));
assert_eq!(128, parse_image_size(Some("-1")));
assert_eq!(128, parse_image_size(Some("99999999999")));
}
#[test]
-1
View File
@@ -2,7 +2,6 @@
mod aggregate_error;
pub mod asset_hash;
mod asset_size;
pub mod aws_sigv4;
mod byte_budget;
mod byte_cache;
+2 -5
View File
@@ -1,9 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-or-later
use crate::{
asset_size,
constants::{AssetExtension, AssetKind},
};
use crate::constants::{self, AssetExtension, AssetKind};
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum OutputFormat {
@@ -97,7 +94,7 @@ pub fn select_url_variant(input: Input) -> OutputSelection {
format: OutputFormat::coerce_from_extension(requested),
size: input
.requested_size
.map(|size| asset_size::clamp_size(size, input.kind)),
.map(|size| constants::clamp_size(size, input.kind)),
reason: if is_output_format_supported(requested) {
"url"
} else {
+1 -2
View File
@@ -6,7 +6,6 @@ pub(in crate::server) mod response;
mod tests;
use crate::{
asset_size,
coalescer::CoalescerError,
config::PolicyMode,
constants::{self, AssetExtension, AssetKind},
@@ -65,7 +64,7 @@ pub(in crate::server) async fn serve_asset_image(
let range_header = headers.get(header::RANGE).and_then(|v| v.to_str().ok());
let requested_download = bool_param(params, "download", false);
let asset_filename = asset_filename_hint(&asset);
let size = asset_size::parse_image_size(params.get("size").map(String::as_str));
let size = constants::parse_image_size(params.get("size").map(String::as_str));
let requested_manual_format = asset_manual_format_override(params, asset.original_ext);
let manual_format_override = asset.forced_output_format.or(requested_manual_format);
let selected = output_format::select_url_variant(output_format::Input {
@@ -70,7 +70,7 @@ fn test_transform_cache() -> crate::transform_cache::TransformCache {
#[test]
fn asset_size_query_is_clamped_by_kind() {
let params = HashMap::from([("size".to_owned(), "4096".to_owned())]);
let size = asset_size::parse_image_size(params.get("size").map(String::as_str));
let size = constants::parse_image_size(params.get("size").map(String::as_str));
let selected = output_format::select_url_variant(output_format::Input {
kind: AssetKind::Avatar,
original: AssetExtension::Webp,
@@ -339,7 +339,7 @@ async fn an_asset_image_source_that_vanishes_mid_request_is_not_found() {
}
fn avatar_cache_key_for_requested_size(raw: &str) -> String {
let size = asset_size::parse_image_size(Some(raw));
let size = constants::parse_image_size(Some(raw));
let selected = output_format::select_url_variant(output_format::Input {
kind: AssetKind::Avatar,
original: AssetExtension::Webp,
-28
View File
@@ -26,31 +26,3 @@ pub fn init_tracing() {
.try_init()
.ok();
}
pub async fn run<S>(service: S) -> anyhow::Result<()>
where
S: router::RouterService,
S::Response: Clone,
{
init_tracing();
let config = config::ServiceConfig::from_env()?;
let transport =
transport::NatsTransport::connect(&config.nats_url, config.nats_auth_token.as_deref())
.await?;
tracing::info!(
service = config.service_name,
mode = ?config.mode,
shard_id = config.shard_id,
shard_count = config.shard_count,
listen_addr = %config.listen_addr,
"starting service"
);
match config.mode {
config::Mode::Router => router::run_router(&config, service, transport).await,
config::Mode::Shard => {
anyhow::bail!(
"run() with RouterService cannot be used in shard mode; use run_shard() directly"
)
}
}
}
+2 -7
View File
@@ -6,7 +6,7 @@ use crate::common::{
remove_dir_if_exists, require_env, resolve_calver, run_command, runner_temp, s3_client,
trim_option, upload_s3_plan_append_only,
};
use crate::functions::sha256_reader;
use crate::functions::sha256_file;
use anyhow::{Context, Result, anyhow, ensure};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
@@ -15,7 +15,7 @@ use clap::{Args, ValueEnum};
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::fs::{self, File};
use std::fs;
use std::path::{Path, PathBuf};
const DEFAULT_PUBLIC_ASSET_BASE_URL: &str = "https://fluxerstatic.com";
@@ -497,11 +497,6 @@ fn asset_tree_digests(root: &Path) -> Result<BTreeMap<String, String>> {
.collect()
}
fn sha256_file(path: &Path) -> Result<String> {
let file = File::open(path).with_context(|| format!("Failed to open {}", path.display()))?;
sha256_reader(file).with_context(|| format!("Failed to read {}", path.display()))
}
fn tree_differences(
canonical: &BTreeMap<String, String>,
other: &BTreeMap<String, String>,