fix(gateway): keep dispatch ordered under relay backpressure (#2300)

fluxer_gateway/src/gateway/gateway_dispatch_relay.erl

@@ -41,9 +41,7 @@ start_link() ->
 
 -spec dispatch(pid(), atom(), term()) -> ok.
 dispatch(SessionPid, Event, Payload) ->
-    gateway_dispatch_relay_batch:relay_or_direct(
-        SessionPid, Event, Payload, gateway_dispatch_relay_batch:max_queue()
-    ).
+    gateway_dispatch_relay_batch:relay_or_direct(SessionPid, Event, Payload).
 
 -spec dispatch(pid(), atom(), term(), term()) -> ok.
 dispatch(SessionPid, Event, Payload, _IgnoredPartitionKey) ->
@@ -51,9 +49,7 @@ dispatch(SessionPid, Event, Payload, _IgnoredPartitionKey) ->
 
 -spec dispatch_many([pid()], atom(), term()) -> ok.
 dispatch_many(SessionPids, Event, Payload) ->
-    gateway_dispatch_relay_batch:relay_or_direct_many(
-        SessionPids, Event, Payload, gateway_dispatch_relay_batch:max_queue()
-    ).
+    gateway_dispatch_relay_batch:relay_or_direct_many(SessionPids, Event, Payload).
 
 -spec dispatch_many([pid()], atom(), term(), term()) -> ok.
 dispatch_many(SessionPids, Event, Payload, _IgnoredPartitionKey) ->

fluxer_gateway/src/gateway/gateway_dispatch_relay_batch.erl

@@ -4,9 +4,8 @@
 -typing([eqwalizer]).
 
 -export([
-    relay_or_direct_many/4,
-    relay_or_direct/4,
-    max_queue/0,
+    relay_or_direct_many/3,
+    relay_or_direct/3,
     select_worker/1,
     current_workers/0,
     current_workers_tuple/0,
@@ -20,8 +19,8 @@
 
 -define(STATE_KEY, {gateway_dispatch_relay, state}).
 
--spec relay_or_direct_many([pid()], atom(), term(), non_neg_integer()) -> ok.
-relay_or_direct_many(SessionPids, Event, Payload, MaxQueue) ->
+-spec relay_or_direct_many([pid()], atom(), term()) -> ok.
+relay_or_direct_many(SessionPids, Event, Payload) ->
     Workers = current_workers_tuple_normalized(),
     case tuple_size(Workers) of
         0 ->
@@ -33,14 +32,13 @@ relay_or_direct_many(SessionPids, Event, Payload, MaxQueue) ->
                 fun gateway_dispatch_relay:dispatch_direct/3
             );
         Count ->
-            relay_many_to_shards(SessionPids, Event, Payload, Workers, Count, MaxQueue)
+            relay_many_to_shards(SessionPids, Event, Payload, Workers, Count)
     end.
 
--spec relay_many_to_shards([pid()], atom(), term(), tuple(), pos_integer(), non_neg_integer()) ->
-    ok.
-relay_many_to_shards(SessionPids, Event, Payload, Workers, Count, MaxQueue) ->
+-spec relay_many_to_shards([pid()], atom(), term(), tuple(), pos_integer()) -> ok.
+relay_many_to_shards(SessionPids, Event, Payload, Workers, Count) ->
     ShardBuckets = build_shard_buckets(SessionPids, Count),
-    deliver_shard_buckets(1, Count, ShardBuckets, Event, Payload, Workers, MaxQueue).
+    deliver_shard_buckets(1, Count, ShardBuckets, Event, Payload, Workers).
 
 -spec build_shard_buckets([pid()], pos_integer()) -> tuple().
 build_shard_buckets(SessionPids, Count) ->
@@ -57,63 +55,30 @@ build_shard_buckets(SessionPids, Count) ->
     ).
 
 -spec deliver_shard_buckets(
-    pos_integer(), pos_integer(), tuple(), atom(), term(), tuple(), non_neg_integer()
+    pos_integer(), pos_integer(), tuple(), atom(), term(), tuple()
 ) -> ok.
-deliver_shard_buckets(Index, Count, _Buckets, _Event, _Payload, _Workers, _MaxQueue) when
-    Index > Count
-->
+deliver_shard_buckets(Index, Count, _Buckets, _Event, _Payload, _Workers) when Index > Count ->
     ok;
-deliver_shard_buckets(Index, Count, Buckets, Event, Payload, Workers, MaxQueue) ->
+deliver_shard_buckets(Index, Count, Buckets, Event, Payload, Workers) ->
     case element(Index, Buckets) of
         [] -> ok;
-        Pids -> deliver_shard(Index, Pids, Event, Payload, Workers, MaxQueue)
+        Pids -> deliver_shard(Index, Pids, Event, Payload, Workers)
     end,
-    deliver_shard_buckets(Index + 1, Count, Buckets, Event, Payload, Workers, MaxQueue).
+    deliver_shard_buckets(Index + 1, Count, Buckets, Event, Payload, Workers).
 
--spec deliver_shard(pos_integer(), [pid()], atom(), term(), tuple(), non_neg_integer()) -> ok.
-deliver_shard(Index, Pids, Event, Payload, Workers, MaxQueue) ->
-    Worker = element(Index, Workers),
-    case is_backpressured(Worker, MaxQueue) of
-        true ->
-            Grouped = gateway_dispatch_relay:group_by_node(Pids),
-            gateway_dispatch_relay:dispatch_grouped(
-                Grouped,
-                Event,
-                Payload,
-                fun gateway_dispatch_relay:dispatch_direct/3
-            );
-        false ->
-            gen_server:cast(Worker, {deliver_many, Pids, Event, Payload})
-    end,
-    ok.
+-spec deliver_shard(pos_integer(), [pid()], atom(), term(), tuple()) -> ok.
+deliver_shard(Index, Pids, Event, Payload, Workers) ->
+    gen_server:cast(element(Index, Workers), {deliver_many, Pids, Event, Payload}).
 
--spec relay_or_direct(pid(), atom(), term(), non_neg_integer()) -> ok.
-relay_or_direct(SessionPid, Event, Payload, MaxQueue) ->
+-spec relay_or_direct(pid(), atom(), term()) -> ok.
+relay_or_direct(SessionPid, Event, Payload) ->
     case select_worker(SessionPid) of
         undefined ->
             gateway_dispatch_relay:dispatch_direct(SessionPid, Event, Payload);
         Worker ->
-            relay_to_worker(Worker, SessionPid, Event, Payload, MaxQueue)
+            gen_server:cast(Worker, {deliver, SessionPid, Event, Payload})
     end.
 
--spec relay_to_worker(pid(), pid(), atom(), term(), non_neg_integer()) -> ok.
-relay_to_worker(Worker, SessionPid, Event, Payload, MaxQueue) ->
-    case is_backpressured(Worker, MaxQueue) of
-        true ->
-            gateway_dispatch_relay:dispatch_direct(SessionPid, Event, Payload);
-        false ->
-            gen_server:cast(Worker, {deliver, SessionPid, Event, Payload}),
-            ok
-    end.
-
--spec is_backpressured(pid(), non_neg_integer()) -> boolean().
-is_backpressured(Worker, MaxQueue) ->
-    MaxQueue > 0 andalso message_queue_len(Worker) >= MaxQueue.
-
--spec max_queue() -> non_neg_integer().
-max_queue() ->
-    gateway_rollout_config:gateway_dispatch_relay_max_queue().
-
 -spec current_workers() -> [pid()].
 current_workers() ->
     tuple_to_list(current_workers_tuple()).

fluxer_gateway/test/gateway_dispatch_relay_ordering_tests.erl

@@ -0,0 +1,98 @@
+%% 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_rollout_config:gateway_dispatch_relay_max_queue(),
+    ok = sys:suspend(Worker),
+    fill_queue(Worker, MaxQueue - 1),
+    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.

fluxer_gateway/test/gateway_dispatch_relay_stress_tests.erl

@@ -14,8 +14,8 @@ dispatch_many_direct_fallback_delivers_once_to_many_sessions_test_() ->
 dispatch_many_worker_shards_deliver_once_to_many_sessions_test_() ->
     {timeout, 30, fun dispatch_many_worker_shards_deliver_once_to_many_sessions/0}.
 
-dispatch_many_backpressured_worker_falls_back_to_direct_delivery_test_() ->
-    {timeout, 30, fun dispatch_many_backpressured_worker_falls_back_to_direct_delivery/0}.
+dispatch_many_backpressured_worker_still_relays_delivery_test_() ->
+    {timeout, 30, fun dispatch_many_backpressured_worker_still_relays_delivery/0}.
 
 dispatch_many_direct_fallback_delivers_once_to_many_sessions() ->
     with_relay_state_cleared(fun() ->
@@ -44,24 +44,26 @@ dispatch_many_worker_shards_deliver_once_to_many_sessions() ->
     end),
     stop_workers(Workers).
 
-dispatch_many_backpressured_worker_falls_back_to_direct_delivery() ->
-    BackpressuredWorker = start_backpressured_worker(64),
-    with_relay_workers([BackpressuredWorker], fun() ->
-        ok = wait_until(fun() ->
-            gateway_dispatch_relay_batch:message_queue_len(BackpressuredWorker) >= 64
-        end),
+dispatch_many_backpressured_worker_still_relays_delivery() ->
+    Worker = gateway_dispatch_relay_batch:start_worker(0),
+    with_relay_workers([Worker], fun() ->
+        ok = sys:suspend(Worker),
+        backpressure_worker(Worker, 64),
         {Receivers, Ref} = start_receivers(512),
         try
             Payload = #{<<"stress">> => <<"backpressure">>},
             ok = gateway_dispatch_relay_batch:relay_or_direct_many(
-                Receivers, relay_stress_event, Payload, 1
+                Receivers, relay_stress_event, Payload
             ),
-            assert_received_once(Receivers, Ref, relay_stress_event, Payload)
+            ?assert(gateway_dispatch_relay_batch:message_queue_len(Worker) >= 64),
+            ok = sys:resume(Worker),
+            assert_received_once(Receivers, Ref, relay_stress_event, Payload),
+            assert_worker_queues_drained([Worker])
         after
             stop_receivers(Receivers)
         end
     end),
-    BackpressuredWorker ! stop.
+    stop_worker(Worker).
 
 with_relay_state_cleared(Fun) ->
     Previous = persistent_term:get(?STATE_KEY, undefined),
@@ -124,17 +126,8 @@ stop_worker(Pid) ->
         exit:_ -> ok
     end.
 
-start_backpressured_worker(MessageCount) ->
-    Pid = spawn_link(fun backpressured_worker_loop/0),
-    lists:foreach(fun(I) -> Pid ! {queued, I} end, lists:seq(1, MessageCount)),
-    Pid.
-
-backpressured_worker_loop() ->
-    receive
-        stop -> ok
-    after 30000 ->
-        ok
-    end.
+backpressure_worker(Worker, MessageCount) ->
+    lists:foreach(fun(I) -> Worker ! {queued, I} end, lists:seq(1, MessageCount)).
 
 assert_received_once(Receivers, Ref, Event, Payload) ->
     Expected = maps:from_list([{Pid, false} || Pid <- Receivers]),