fix(gateway): wire presence grace, request limits and voice sync (#3329)

fluxer_gateway/test/gateway_cluster_membership_tests.erl

@@ -291,6 +291,51 @@ update_members_role_lookups_run_concurrently_test() ->
         ?assertEqual(lists:usort(Nodes), maps:get(sessions, maps:get(members_by_role, State1)))
     end).
 
+member_nodedown_starts_presence_grace_and_nodeup_cancels_it_test() ->
+    with_clean_cluster_terms(fun() ->
+        with_presence_cache(fun(CachePid) ->
+            Peer = 'peer@a',
+            State0 = (base_state(#{members => lists:usort([node(), Peer])}))#{
+                discovered => [Peer]
+            },
+            {noreply, State1} = gateway_cluster_membership:handle_info(
+                {nodedown, Peer}, State0
+            ),
+            ?assertNot(lists:member(Peer, maps:get(members, State1))),
+            ?assert(maps:is_key(Peer, pending_nodedown_cleanups(CachePid))),
+            {noreply, State2} = gateway_cluster_membership:handle_info({nodeup, Peer}, State1),
+            ?assert(lists:member(Peer, maps:get(members, State2))),
+            ?assertEqual(#{}, pending_nodedown_cleanups(CachePid))
+        end)
+    end).
+
+non_member_nodedown_does_not_start_presence_grace_test() ->
+    with_clean_cluster_terms(fun() ->
+        with_presence_cache(fun(CachePid) ->
+            State0 = base_state(#{}),
+            {noreply, _State1} = gateway_cluster_membership:handle_info(
+                {nodedown, 'remsh@host'}, State0
+            ),
+            ?assertEqual(#{}, pending_nodedown_cleanups(CachePid))
+        end)
+    end).
+
+with_presence_cache(Fun) ->
+    {ok, Pid} =
+        case whereis(presence_cache) of
+            undefined -> presence_cache:start_link();
+            Existing -> {ok, Existing}
+        end,
+    try
+        Fun(Pid)
+    after
+        unlink(Pid),
+        gen_server:stop(Pid)
+    end.
+
+pending_nodedown_cleanups(CachePid) ->
+    maps:get(pending_nodedown_cleanups, sys:get_state(CachePid)).
+
 base_state(Overrides) ->
     Defaults = #{
         members => [node()],

fluxer_gateway/test/guild_request_members_filter_tests.erl

@@ -210,6 +210,143 @@ check_guild_request_rate_limit_invalid_guild_test() ->
         guild_request_members_filter:check_guild_request_rate_limit(invalid_guild_id())
     ).
 
+handle_request_consumes_user_budget_test() ->
+    with_rate_limits_enabled(fun() ->
+        UserId = 555000001,
+        clear_request_rate_limit(UserId),
+        SessionState = #{user_id => integer_to_binary(UserId), guilds => #{}, bot => false},
+        Data = #{<<"guild_id">> => <<"777000001">>, <<"query">> => <<"a">>, <<"limit">> => 10},
+        Results = [
+            guild_request_members:handle_request(Data, self(), SessionState)
+         || _ <- lists:seq(1, ?REQUEST_MEMBERS_RATE_LIMIT_MAX_EVENTS)
+        ],
+        ?assertEqual(
+            lists:duplicate(?REQUEST_MEMBERS_RATE_LIMIT_MAX_EVENTS, {error, guild_not_found}),
+            Results
+        ),
+        ?assertEqual(
+            {error, rate_limited},
+            guild_request_members:handle_request(Data, self(), SessionState)
+        ),
+        clear_request_rate_limit(UserId)
+    end).
+
+handle_request_refuses_guild_over_budget_test() ->
+    with_rate_limits_enabled(fun() ->
+        UserId = 555000002,
+        GuildIds = [777000002, 777000003],
+        SessionState = spent_guild_session(UserId, GuildIds),
+        [GuildIdBin | _] = GuildIdBins = [integer_to_binary(Id) || Id <- GuildIds],
+        ?assertEqual(
+            {error, rate_limited},
+            guild_request_members:handle_request(
+                #{<<"guild_id">> => GuildIdBin, <<"query">> => <<"a">>}, self(), SessionState
+            )
+        ),
+        ?assertEqual(
+            {error, guild_not_found},
+            guild_request_members:handle_request(
+                #{<<"guild_ids">> => GuildIdBins, <<"query">> => <<"a">>}, self(), SessionState
+            )
+        ),
+        clear_request_rate_limit(UserId),
+        lists:foreach(fun clear_guild_request_rate_limit/1, GuildIds)
+    end).
+
+spent_guild_session(UserId, GuildIds) ->
+    clear_request_rate_limit(UserId),
+    ensure_ets_table(?REQUEST_MEMBERS_GUILD_RATE_LIMIT_TABLE),
+    Now = erlang:system_time(millisecond),
+    Spent = lists:duplicate(?REQUEST_MEMBERS_GUILD_RATE_LIMIT_MAX_EVENTS, Now - 1000),
+    ets:insert(?REQUEST_MEMBERS_GUILD_RATE_LIMIT_TABLE, [{Id, Spent} || Id <- GuildIds]),
+    #{
+        user_id => integer_to_binary(UserId),
+        guilds => maps:from_list([{Id, {self(), make_ref()}} || Id <- GuildIds]),
+        bot => false
+    }.
+
+handle_request_budgets_disabled_by_env_test() ->
+    OldValue = os:getenv("FLUXER_DISABLE_RATE_LIMITS"),
+    os:putenv("FLUXER_DISABLE_RATE_LIMITS", "true"),
+    try
+        UserId = 555000004,
+        clear_request_rate_limit(UserId),
+        SessionState = #{user_id => integer_to_binary(UserId), guilds => #{}, bot => false},
+        Data = #{<<"guild_id">> => <<"777000004">>, <<"query">> => <<"a">>},
+        Results = [
+            guild_request_members:handle_request(Data, self(), SessionState)
+         || _ <- lists:seq(1, ?REQUEST_MEMBERS_RATE_LIMIT_MAX_EVENTS + 1)
+        ],
+        ?assertEqual(
+            lists:duplicate(
+                ?REQUEST_MEMBERS_RATE_LIMIT_MAX_EVENTS + 1, {error, guild_not_found}
+            ),
+            Results
+        )
+    after
+        restore_env("FLUXER_DISABLE_RATE_LIMITS", OldValue)
+    end.
+
+rate_limit_tables_outlive_the_request_worker_test() ->
+    with_rate_limits_enabled(fun() ->
+        UserId = 555000005,
+        drop_table(?REQUEST_MEMBERS_RATE_LIMIT_TABLE),
+        {ok, OwnerPid} = start_ets_owner(),
+        try
+            ok = check_rate_limit_in_worker(UserId),
+            ?assertMatch(
+                [{UserId, [_]}], ets:lookup(?REQUEST_MEMBERS_RATE_LIMIT_TABLE, UserId)
+            ),
+            ?assertEqual(OwnerPid, ets:info(?REQUEST_MEMBERS_RATE_LIMIT_TABLE, owner))
+        after
+            gen_server:stop(OwnerPid),
+            drop_table(?REQUEST_MEMBERS_RATE_LIMIT_TABLE)
+        end
+    end).
+
+check_rate_limit_in_worker(UserId) ->
+    Self = self(),
+    {Worker, Ref} = spawn_monitor(fun() ->
+        Self ! {checked, guild_request_members_filter:check_request_rate_limit(UserId)}
+    end),
+    receive
+        {checked, ok} -> ok
+    after 1000 -> error(worker_timeout)
+    end,
+    receive
+        {'DOWN', Ref, process, Worker, _} -> ok
+    after 1000 -> error(worker_alive)
+    end.
+
+drop_table(Name) ->
+    case ets:whereis(Name) of
+        undefined -> ok;
+        _ -> ets:delete(Name)
+    end.
+
+prune_expired_entries_drops_idle_keys_test() ->
+    Table = ets:new(request_members_prune_test, [public, set]),
+    Now = erlang:system_time(millisecond),
+    ets:insert(Table, [
+        {idle_window, [Now - 20000, Now - 25000]},
+        {active_window, [Now - 100, Now - 20000]},
+        {idle_bot, Now - 40000},
+        {active_bot, Now - 100}
+    ]),
+    ?assertEqual(ok, guild_request_members_filter:prune_expired_entries(Table, 10000)),
+    ?assertEqual(
+        [active_bot, active_window],
+        lists:sort([Key || {Key, _} <- ets:tab2list(Table)])
+    ),
+    ets:delete(Table),
+    ?assertEqual(gone, guild_request_members_filter:prune_expired_entries(Table, 10000)).
+
+start_ets_owner() ->
+    case whereis(guild_ets_owner) of
+        undefined -> guild_ets_owner:start_link();
+        Existing -> {ok, Existing}
+    end.
+
 restore_env(Key, false) ->
     os:unsetenv(Key);
 restore_env(Key, Value) ->

fluxer_gateway/test/guild_sessions_connect_tests.erl

@@ -141,6 +141,72 @@ normalize_connect_queue_test() ->
     ?assertEqual(Q, guild_sessions_connect_cleanup:normalize_connect_queue(Q)),
     ?assertEqual(undefined, guild_sessions_connect_cleanup:normalize_connect_queue(undefined)).
 
+session_down_clears_passive_sync_entry_test() ->
+    ok = passive_sync_registry:init(),
+    SessionId = <<"passive-down">>,
+    ok = passive_sync_registry:store(SessionId, 42, stored_passive_state()),
+    {Ref, Session, State} = down_state(SessionId),
+    {noreply, _} = guild_sessions_connect:handle_session_down(Ref, {SessionId, Session}, State),
+    ?assertEqual([], ets:lookup(passive_sync_registry, {SessionId, 42})).
+
+remove_session_clears_passive_sync_entry_test() ->
+    ok = passive_sync_registry:init(),
+    SessionId = <<"passive-remove">>,
+    ok = passive_sync_registry:store(SessionId, 42, stored_passive_state()),
+    _ = guild_sessions_connect:remove_session(SessionId, remove_session_state(SessionId)),
+    ?assertEqual([], ets:lookup(passive_sync_registry, {SessionId, 42})).
+
+passive_sync_registry_stays_bounded_across_session_cycles_test() ->
+    ok = passive_sync_registry:init(),
+    GuildRows = fun() -> length(ets:match(passive_sync_registry, {{'_', 42}, '_'})) end,
+    Before = GuildRows(),
+    lists:foreach(
+        fun(N) ->
+            SessionId = <<"cycle-", (integer_to_binary(N))/binary>>,
+            ok = passive_sync_registry:store(SessionId, 42, stored_passive_state()),
+            {Ref, Session, State} = down_state(SessionId),
+            {noreply, _} = guild_sessions_connect:handle_session_down(
+                Ref, {SessionId, Session}, State
+            )
+        end,
+        lists:seq(1, 500)
+    ),
+    ?assertEqual(Before, GuildRows()).
+
+guild_terminate_clears_passive_sync_entries_test() ->
+    ok = passive_sync_registry:init(),
+    GuildId = 4242,
+    OtherGuildId = 4343,
+    lists:foreach(
+        fun({SessionId, GId}) ->
+            ok = passive_sync_registry:store(SessionId, GId, stored_passive_state())
+        end,
+        [{<<"t1">>, GuildId}, {<<"t2">>, GuildId}, {<<"t1">>, OtherGuildId}]
+    ),
+    ok = guild:terminate(shutdown, #{
+        id => GuildId,
+        sessions => #{<<"t1">> => #{}, <<"t2">> => #{}}
+    }),
+    ?assertEqual([], ets:lookup(passive_sync_registry, {<<"t1">>, GuildId})),
+    ?assertEqual([], ets:lookup(passive_sync_registry, {<<"t2">>, GuildId})),
+    ?assertEqual(
+        stored_passive_state(), passive_sync_registry:lookup(<<"t1">>, OtherGuildId)
+    ),
+    ok = passive_sync_registry:delete(<<"t1">>, OtherGuildId).
+
+stored_passive_state() ->
+    #{
+        previous_passive_updates => #{<<"1">> => <<"2">>},
+        previous_passive_channel_versions => #{},
+        previous_passive_voice_states => #{}
+    }.
+
+down_state(SessionId) ->
+    State0 = remove_session_state(SessionId),
+    Session = maps:get(SessionId, maps:get(sessions, State0)),
+    Ref = maps:get(mref, Session),
+    {Ref, Session, State0#{guild_session_refs => #{Ref => SessionId}}}.
+
 remove_session_state(SessionId) ->
     SessionData = #{
         session_id => SessionId,

fluxer_gateway/test/presence_cache_tests.erl

@@ -90,6 +90,14 @@ nodeup_cancels_grace_period_test() ->
     timer:sleep(50),
     ?assertEqual(ok, gen_server:stop(Pid)).
 
+nodedown_grace_expiry_clears_pending_cleanup_test() ->
+    {ok, Pid} = maybe_start_for_test(),
+    ok = presence_cache:handle_nodedown('lost@node'),
+    ?assertMatch(#{'lost@node' := _}, maps:get(pending_nodedown_cleanups, sys:get_state(Pid))),
+    Pid ! {nodedown_grace_expired, 'lost@node'},
+    ?assertEqual(#{}, maps:get(pending_nodedown_cleanups, sys:get_state(Pid))),
+    ?assertEqual(ok, gen_server:stop(Pid)).
+
 anti_entropy_no_op_when_in_sync_test() ->
     {ok, Pid} = maybe_start_for_test(),
     State = cache_state(sys:get_state(Pid)),