fix(gateway): reclaim ip connection counts from dead sockets (#2308)

fluxer_gateway/src/gateway/gateway_handler.erl

@@ -81,7 +81,7 @@ init(Req, _Opts) ->
 
 -spec websocket_init(state()) -> ws_result().
 websocket_init(#{version := 1} = State) ->
-    case gateway_handler_rate_limit:acquire_connection(maps:get(peer_ip, State, undefined)) of
+    case gateway_ip_connections:acquire(maps:get(peer_ip, State, undefined)) of
         ok ->
             do_websocket_init(State#{connection_acquired => true});
         {error, too_many_connections} ->
@@ -168,7 +168,7 @@ terminate_with_state(_) ->
 
 -spec maybe_release_connection(state()) -> ok.
 maybe_release_connection(#{connection_acquired := true} = State) ->
-    gateway_handler_rate_limit:note_disconnect(State);
+    gateway_ip_connections:note_disconnect(State);
 maybe_release_connection(_) ->
     ok.
 

fluxer_gateway/src/gateway/gateway_handler_rate_limit.erl

@@ -5,11 +5,9 @@
 
 -export([
     check_rate_limit/2,
-    acquire_connection/1,
-    release_connection/1,
     check_shared_ip_rate/1,
     check_shared_user_rate/1,
-    note_disconnect/1
+    rate_limits_disabled/0
 ]).
 
 -export_type([state/0]).
@@ -21,13 +19,11 @@
 
 -define(SHARED_IP_RATE_TABLE, gateway_shared_ip_rate).
 -define(SHARED_USER_RATE_TABLE, gateway_shared_user_rate).
--define(IP_CONNECTION_TABLE, gateway_ip_connections).
 
 -define(SHARED_IP_RATE_WINDOW_MS, 60000).
 -define(SHARED_IP_RATE_MAX_EVENTS, 6000).
 -define(SHARED_USER_RATE_WINDOW_MS, 60000).
 -define(SHARED_USER_RATE_MAX_EVENTS, 600).
--define(MAX_CONNECTIONS_PER_IP, 256).
 -define(SHARED_RATE_CLEANUP_INTERVAL_MS, ?SHARED_IP_RATE_WINDOW_MS * 2).
 
 -type state() :: gateway_handler:state().
@@ -201,63 +197,6 @@ check_shared_window(Table, Key, WindowMs, MaxEvents, LimitReason) ->
         error:badarg -> ok
     end.
 
--spec acquire_connection(term()) -> ok | {error, too_many_connections}.
-acquire_connection(PeerIP) when is_binary(PeerIP), PeerIP =/= <<"unknown">> ->
-    case rate_limits_disabled() of
-        true ->
-            ok;
-        false ->
-            ensure_counter_table(?IP_CONNECTION_TABLE),
-            do_acquire_connection(PeerIP)
-    end;
-acquire_connection(_) ->
-    ok.
-
--spec do_acquire_connection(binary()) -> ok | {error, too_many_connections}.
-do_acquire_connection(PeerIP) ->
-    try ets:update_counter(?IP_CONNECTION_TABLE, PeerIP, {2, 1}, {PeerIP, 0}) of
-        Count when Count > ?MAX_CONNECTIONS_PER_IP ->
-            _ = ets:update_counter(?IP_CONNECTION_TABLE, PeerIP, {2, -1, 0, 0}),
-            {error, too_many_connections};
-        _ ->
-            ok
-    catch
-        error:badarg -> ok
-    end.
-
--spec release_connection(term()) -> ok.
-release_connection(PeerIP) when is_binary(PeerIP), PeerIP =/= <<"unknown">> ->
-    case ets:whereis(?IP_CONNECTION_TABLE) of
-        undefined ->
-            ok;
-        _ ->
-            decrement_connection(PeerIP)
-    end;
-release_connection(_) ->
-    ok.
-
--spec decrement_connection(binary()) -> ok.
-decrement_connection(PeerIP) ->
-    try update_connection_count(PeerIP) of
-        ok -> ok
-    catch
-        error:badarg -> ok
-    end.
-
--spec update_connection_count(binary()) -> ok.
-update_connection_count(PeerIP) ->
-    case ets:update_counter(?IP_CONNECTION_TABLE, PeerIP, {2, -1, 0, 0}) of
-        0 ->
-            ets:delete(?IP_CONNECTION_TABLE, PeerIP),
-            ok;
-        _ ->
-            ok
-    end.
-
--spec note_disconnect(state()) -> ok.
-note_disconnect(State) ->
-    release_connection(maps:get(peer_ip, State, undefined)).
-
 -spec ensure_window_table(atom(), pos_integer()) -> ok.
 ensure_window_table(Table, WindowMs) ->
     case ets:whereis(Table) of
@@ -305,18 +244,6 @@ prune_old_window_entries(Table, WindowMs) ->
         error:badarg -> gone
     end.
 
--spec ensure_counter_table(atom()) -> ok.
-ensure_counter_table(Table) ->
-    case ets:whereis(Table) of
-        undefined -> create_counter_table(Table);
-        _ -> ok
-    end.
-
--spec create_counter_table(atom()) -> ok.
-create_counter_table(Table) ->
-    _ = create_table(Table),
-    ok.
-
 -spec create_table(atom()) -> created | exists.
 create_table(Table) ->
     try
@@ -388,38 +315,12 @@ shared_ip_rate_ignores_unknown_ip_test() ->
         ?assertEqual(ok, check_shared_ip_rate(undefined))
     end).
 
-connection_cap_blocks_over_limit_test() ->
-    with_rate_limits_enabled(fun() ->
-        IP = <<"198.51.100.20">>,
-        reset_connections(IP),
-        assert_connection_cap_allows_limit(IP),
-        ?assertEqual({error, too_many_connections}, acquire_connection(IP)),
-        ok = release_connection(IP),
-        ?assertEqual(ok, acquire_connection(IP)),
-        reset_connections(IP)
-    end).
-
 assert_shared_ip_rate_allows_limit(IP) ->
     lists:foreach(
         fun(_) -> ?assertEqual(ok, check_shared_ip_rate(IP)) end,
         lists:seq(1, ?SHARED_IP_RATE_MAX_EVENTS)
     ).
 
-assert_connection_cap_allows_limit(IP) ->
-    lists:foreach(
-        fun(_) -> ?assertEqual(ok, acquire_connection(IP)) end,
-        lists:seq(1, ?MAX_CONNECTIONS_PER_IP)
-    ).
-
-connection_release_decrements_test() ->
-    with_rate_limits_enabled(fun() ->
-        IP = <<"198.51.100.30">>,
-        reset_connections(IP),
-        ok = acquire_connection(IP),
-        ok = release_connection(IP),
-        ?assertEqual([], ets:lookup(?IP_CONNECTION_TABLE, IP))
-    end).
-
 reset_shared_ip(IP) ->
     case ets:whereis(?SHARED_IP_RATE_TABLE) of
         undefined ->
@@ -606,13 +507,4 @@ await_window_cleanup_loops(Max, Attempts) ->
             await_window_cleanup_loops(Max, Attempts - 1)
     end.
 
-reset_connections(IP) ->
-    case ets:whereis(?IP_CONNECTION_TABLE) of
-        undefined ->
-            ok;
-        _ ->
-            ets:delete(?IP_CONNECTION_TABLE, IP),
-            ok
-    end.
-
 -endif.

fluxer_gateway/src/gateway/gateway_ip_connections.erl

@@ -0,0 +1,347 @@
+%% SPDX-License-Identifier: AGPL-3.0-or-later
+
+-module(gateway_ip_connections).
+-typing([eqwalizer]).
+
+-export([acquire/1, release/1, note_disconnect/1]).
+
+-define(CONNECTION_TABLE, gateway_ip_connections).
+-define(OWNER_TABLE, gateway_ip_connection_owners).
+-define(MAX_CONNECTIONS_PER_IP, 256).
+-define(OWNER_SWEEP_INTERVAL_MS, 60000).
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+-endif.
+
+-spec acquire(term()) -> ok | {error, too_many_connections}.
+acquire(PeerIP) when is_binary(PeerIP), PeerIP =/= <<"unknown">> ->
+    case gateway_handler_rate_limit:rate_limits_disabled() of
+        true ->
+            ok;
+        false ->
+            ensure_tables(),
+            do_acquire(PeerIP)
+    end;
+acquire(_) ->
+    ok.
+
+-spec release(term()) -> ok.
+release(PeerIP) when is_binary(PeerIP), PeerIP =/= <<"unknown">> ->
+    case take_owner({self(), PeerIP}) of
+        true -> decrement(PeerIP);
+        false -> ok
+    end;
+release(_) ->
+    ok.
+
+-spec note_disconnect(gateway_handler:state()) -> ok.
+note_disconnect(State) ->
+    release(maps:get(peer_ip, State, undefined)).
+
+-spec do_acquire(binary()) -> ok | {error, too_many_connections}.
+do_acquire(PeerIP) ->
+    try ets:update_counter(?CONNECTION_TABLE, PeerIP, {2, 1}, {PeerIP, 0}) of
+        Count when Count > ?MAX_CONNECTIONS_PER_IP ->
+            _ = ets:update_counter(?CONNECTION_TABLE, PeerIP, {2, -1, 0, 0}),
+            {error, too_many_connections};
+        _ ->
+            track_owner(PeerIP)
+    catch
+        error:badarg -> ok
+    end.
+
+-spec track_owner(binary()) -> ok.
+track_owner(PeerIP) ->
+    try
+        _ = ets:insert(?OWNER_TABLE, {{self(), PeerIP}}),
+        ok
+    catch
+        error:badarg -> ok
+    end.
+
+-spec take_owner({pid(), binary()}) -> boolean().
+take_owner(Owner) ->
+    try ets:take(?OWNER_TABLE, Owner) of
+        [_ | _] -> true;
+        [] -> false
+    catch
+        error:badarg -> false
+    end.
+
+-spec decrement(binary()) -> ok.
+decrement(PeerIP) ->
+    try update_connection_count(PeerIP) of
+        ok -> ok
+    catch
+        error:badarg -> ok
+    end.
+
+-spec update_connection_count(binary()) -> ok.
+update_connection_count(PeerIP) ->
+    case ets:update_counter(?CONNECTION_TABLE, PeerIP, {2, -1, 0, 0}) of
+        0 ->
+            ets:delete(?CONNECTION_TABLE, PeerIP),
+            ok;
+        _ ->
+            ok
+    end.
+
+-spec ensure_tables() -> ok.
+ensure_tables() ->
+    ensure_connection_table(),
+    ensure_owner_table().
+
+-spec ensure_connection_table() -> ok.
+ensure_connection_table() ->
+    case ets:whereis(?CONNECTION_TABLE) of
+        undefined ->
+            _ = create_table(?CONNECTION_TABLE),
+            ok;
+        _ ->
+            ok
+    end.
+
+-spec ensure_owner_table() -> ok.
+ensure_owner_table() ->
+    case ets:whereis(?OWNER_TABLE) of
+        undefined -> create_owner_table();
+        _ -> ok
+    end.
+
+-spec create_owner_table() -> ok.
+create_owner_table() ->
+    case create_table(?OWNER_TABLE) of
+        created -> schedule_owner_sweep();
+        exists -> ok
+    end.
+
+-spec create_table(atom()) -> created | exists.
+create_table(Table) ->
+    try
+        _ = ets:new(Table, table_options()),
+        created
+    catch
+        error:badarg -> exists
+    end.
+
+-spec table_options() -> list().
+table_options() ->
+    [named_table, public, set, {write_concurrency, true}, {read_concurrency, true}] ++
+        guild_ets_utils:heir_options().
+
+-spec schedule_owner_sweep() -> ok.
+schedule_owner_sweep() ->
+    case ets:whereis(?OWNER_TABLE) of
+        undefined -> ok;
+        Tid -> spawn_owner_sweep(Tid)
+    end.
+
+-spec spawn_owner_sweep(ets:table()) -> ok.
+spawn_owner_sweep(Tid) ->
+    _ = spawn(fun() -> owner_sweep_loop(Tid) end),
+    ok.
+
+-spec owner_sweep_loop(ets:table()) -> ok.
+owner_sweep_loop(Table) ->
+    ok = gateway_retry_timer:wait(?OWNER_SWEEP_INTERVAL_MS),
+    case reclaim_dead_owners(Table) of
+        ok -> owner_sweep_loop(Table);
+        gone -> ok
+    end.
+
+-spec reclaim_dead_owners(ets:table()) -> ok | gone.
+reclaim_dead_owners(Table) ->
+    try ets:select(Table, [{{{'$1', '$2'}}, [], [{{'$1', '$2'}}]}]) of
+        Owners -> lists:foreach(fun reclaim_owner/1, Owners)
+    catch
+        error:badarg -> gone
+    end.
+
+-spec reclaim_owner({pid(), binary()}) -> ok.
+reclaim_owner({Pid, PeerIP}) ->
+    case is_process_alive(Pid) of
+        true -> ok;
+        false -> reclaim_dead_owner({Pid, PeerIP}, PeerIP)
+    end.
+
+-spec reclaim_dead_owner({pid(), binary()}, binary()) -> ok.
+reclaim_dead_owner(Owner, PeerIP) ->
+    case take_owner(Owner) of
+        true -> decrement(PeerIP);
+        false -> ok
+    end.
+
+-ifdef(TEST).
+
+with_rate_limits_enabled(Fun) ->
+    OldValue = os:getenv("FLUXER_DISABLE_RATE_LIMITS"),
+    os:unsetenv("FLUXER_DISABLE_RATE_LIMITS"),
+    try
+        Fun()
+    after
+        restore_env("FLUXER_DISABLE_RATE_LIMITS", OldValue)
+    end.
+
+restore_env(Key, false) ->
+    os:unsetenv(Key);
+restore_env(Key, Value) ->
+    os:putenv(Key, Value).
+
+connection_cap_blocks_over_limit_test() ->
+    with_rate_limits_enabled(fun() ->
+        IP = <<"198.51.100.20">>,
+        reset_connections(IP),
+        assert_connection_cap_allows_limit(IP),
+        ?assertEqual({error, too_many_connections}, acquire(IP)),
+        ok = release(IP),
+        ?assertEqual(ok, acquire(IP)),
+        reset_connections(IP)
+    end).
+
+assert_connection_cap_allows_limit(IP) ->
+    lists:foreach(
+        fun(_) -> ?assertEqual(ok, acquire(IP)) end,
+        lists:seq(1, ?MAX_CONNECTIONS_PER_IP)
+    ).
+
+connection_release_decrements_test() ->
+    with_rate_limits_enabled(fun() ->
+        IP = <<"198.51.100.30">>,
+        reset_connections(IP),
+        ok = acquire(IP),
+        ok = release(IP),
+        ?assertEqual([], ets:lookup(?CONNECTION_TABLE, IP))
+    end).
+
+killed_connection_owner_releases_ip_count_test() ->
+    with_rate_limits_enabled(fun() ->
+        with_fast_owner_sweep(fun assert_killed_connection_count_is_released/0)
+    end).
+
+assert_killed_connection_count_is_released() ->
+    with_heir_owner(fun assert_killed_connection_count_drops/0).
+
+assert_killed_connection_count_drops() ->
+    IP = <<"198.51.100.80">>,
+    Holder = start_connection_holder(IP),
+    ?assertEqual(1, connection_count(IP)),
+    kill_connection_holder(Holder),
+    ?assert(await_connection_count(IP, 0, 300)).
+
+released_connection_is_not_released_twice_test() ->
+    with_rate_limits_enabled(fun() ->
+        with_fast_owner_sweep(fun assert_released_connection_is_not_double_released/0)
+    end).
+
+assert_released_connection_is_not_double_released() ->
+    with_heir_owner(fun assert_released_connection_keeps_peer_count/0).
+
+assert_released_connection_keeps_peer_count() ->
+    IP = <<"198.51.100.90">>,
+    Holder = start_connection_holder(IP),
+    Releaser = start_connection_holder(IP),
+    ?assertEqual(2, connection_count(IP)),
+    release_connection_holder(Releaser),
+    ?assert(await_connection_count(IP, 1, 300)),
+    timer:sleep(100),
+    ?assertEqual(1, connection_count(IP)),
+    stop_connection_holder(Holder).
+
+with_fast_owner_sweep(Fun) ->
+    meck:new(gateway_retry_timer, [passthrough]),
+    meck:expect(gateway_retry_timer, wait, fun(_) -> timer:sleep(5) end),
+    try
+        Fun()
+    after
+        meck:unload(gateway_retry_timer)
+    end.
+
+with_heir_owner(Fun) ->
+    drop_table(?CONNECTION_TABLE),
+    drop_table(?OWNER_TABLE),
+    {ok, _} = guild_ets_owner:start_link(),
+    try
+        Fun()
+    after
+        gen_server:stop(guild_ets_owner)
+    end.
+
+drop_table(Table) ->
+    case ets:whereis(Table) of
+        undefined ->
+            ok;
+        _ ->
+            ets:delete(Table),
+            ok
+    end.
+
+start_connection_holder(IP) ->
+    Parent = self(),
+    Pid = spawn(fun() -> hold_connection(Parent, IP) end),
+    receive
+        {connection_acquired, Pid} -> Pid
+    after 1000 -> error(connection_holder_timeout)
+    end.
+
+hold_connection(Parent, IP) ->
+    ok = acquire(IP),
+    Parent ! {connection_acquired, self()},
+    receive
+        release -> release(IP);
+        stop -> ok
+    after 30000 -> ok
+    end.
+
+kill_connection_holder(Pid) ->
+    exit(Pid, kill),
+    await_connection_holder_down(Pid).
+
+release_connection_holder(Pid) ->
+    Pid ! release,
+    await_connection_holder_down(Pid).
+
+stop_connection_holder(Pid) ->
+    Pid ! stop,
+    await_connection_holder_down(Pid).
+
+await_connection_holder_down(Pid) ->
+    Ref = erlang:monitor(process, Pid),
+    receive
+        {'DOWN', Ref, process, Pid, _Reason} -> ok
+    after 1000 -> error(connection_holder_down_timeout)
+    end.
+
+connection_count(IP) ->
+    try ets:lookup(?CONNECTION_TABLE, IP) of
+        [{_, Count}] -> Count;
+        [] -> 0
+    catch
+        error:badarg -> 0
+    end.
+
+await_connection_count(IP, Expected, 0) ->
+    connection_count(IP) =:= Expected;
+await_connection_count(IP, Expected, Attempts) ->
+    case connection_count(IP) =:= Expected of
+        true ->
+            true;
+        false ->
+            timer:sleep(10),
+            await_connection_count(IP, Expected, Attempts - 1)
+    end.
+
+reset_connections(IP) ->
+    _ = take_owner({self(), IP}),
+    drop_connection_row(IP).
+
+drop_connection_row(IP) ->
+    case ets:whereis(?CONNECTION_TABLE) of
+        undefined ->
+            ok;
+        _ ->
+            ets:delete(?CONNECTION_TABLE, IP),
+            ok
+    end.
+
+-endif.