Implement node purge in stream SAC coordinator

acogoluegnes · acogoluegnes · commit 90ba8e5bc9f1 · 2025-04-17T16:10:53.000+02:00
diff --git a/deps/rabbit/src/rabbit_stream_coordinator.erl b/deps/rabbit/src/rabbit_stream_coordinator.erl
@@ -812,7 +812,7 @@ members() ->
             end
     end.
 
-maybe_resize_coordinator_cluster(MachineVersion) ->
+maybe_resize_coordinator_cluster(LeaderPid, MachineVersion) ->
     spawn(fun() ->
                   RabbitIsRunning = rabbit:is_running(),
                   case members() of
@@ -848,27 +848,29 @@ maybe_resize_coordinator_cluster(MachineVersion) ->
                                       remove_member(Leader, Members, Old)
                               end,
                           maybe_handle_stale_nodes(MemberNodes, RabbitNodes,
+                                                   LeaderPid,
                                                    MachineVersion);
-                          _ ->
+                      _ ->
                           ok
                   end
           end).
 
 maybe_handle_stale_nodes(MemberNodes, ExpectedNodes,
-                         MachineVersion) when MachineVersion > 4 ->
+                         LeaderPid, MachineVersion) when MachineVersion > 4 ->
     case MemberNodes -- ExpectedNodes of
         [] ->
             ok;
         Stale when length(ExpectedNodes) > 0 ->
             rabbit_log:debug("Stale nodes detected in stream SAC "
                              "coordinator: ~w. Purging state.",
                              [Stale]),
-            %% TODO SAC pipeline command to purge state from stale nodes
+            Mod = sac_module(MachineVersion),
+            ra:pipeline_command(LeaderPid, Mod:make_purge_nodes(Stale)),
             ok;
         _ ->
             ok
     end;
-maybe_handle_stale_nodes(_, _, _) ->
+maybe_handle_stale_nodes(_, _, _, _) ->
     ok.
 
 add_member(Members, Node) ->
@@ -945,8 +947,9 @@ init_aux(_Name) ->
 %% TODO ensure the dead writer is restarted as a replica at some point in time, increasing timeout?
 handle_aux(leader, _, maybe_resize_coordinator_cluster,
            #aux{resizer = undefined} = Aux, RaAux) ->
+    Leader = ra_aux:leader_id(RaAux),
     MachineVersion = ra_aux:effective_machine_version(RaAux),
-    Pid = maybe_resize_coordinator_cluster(MachineVersion),
+    Pid = maybe_resize_coordinator_cluster(Leader, MachineVersion),
     {no_reply, Aux#aux{resizer = Pid}, RaAux, [{monitor, process, aux, Pid}]};
 handle_aux(leader, _, maybe_resize_coordinator_cluster,
            AuxState, RaAux) ->
diff --git a/deps/rabbit/src/rabbit_stream_sac_coordinator.erl b/deps/rabbit/src/rabbit_stream_sac_coordinator.erl
@@ -21,7 +21,8 @@
 -opaque command() :: #command_register_consumer{} |
                      #command_unregister_consumer{} |
                      #command_activate_consumer{} |
-                     #command_connection_reconnected{}.
+                     #command_connection_reconnected{} |
+                     #command_purge_nodes{}.
 
 -opaque state() :: #?MODULE{}.
 
@@ -47,6 +48,7 @@
          group_consumers/5,
          overview/1,
          import_state/2]).
+-export([make_purge_nodes/1]).
 
 %% exported for unit tests only
 -ifdef(TEST).
@@ -84,25 +86,13 @@ register_consumer(VirtualHost,
                   ConnectionPid,
                   Owner,
                   SubscriptionId) ->
-    process_command({sac,
-                     #command_register_consumer{vhost =
-                                                VirtualHost,
-                                                stream =
-                                                Stream,
-                                                partition_index
-                                                =
-                                                PartitionIndex,
-                                                consumer_name
-                                                =
-                                                ConsumerName,
-                                                connection_pid
-                                                =
-                                                ConnectionPid,
-                                                owner =
-                                                Owner,
-                                                subscription_id
-                                                =
-                                                SubscriptionId}}).
+    process_command(#command_register_consumer{vhost = VirtualHost,
+                                               stream = Stream,
+                                               partition_index = PartitionIndex,
+                                               consumer_name = ConsumerName,
+                                               connection_pid = ConnectionPid,
+                                               owner = Owner,
+                                               subscription_id = SubscriptionId}).
 
 -spec unregister_consumer(binary(),
                           binary(),
@@ -115,35 +105,24 @@ unregister_consumer(VirtualHost,
                     ConsumerName,
                     ConnectionPid,
                     SubscriptionId) ->
-    process_command({sac,
-                     #command_unregister_consumer{vhost =
-                                                  VirtualHost,
-                                                  stream =
-                                                  Stream,
-                                                  consumer_name
-                                                  =
-                                                  ConsumerName,
-                                                  connection_pid
-                                                  =
-                                                  ConnectionPid,
-                                                  subscription_id
-                                                  =
-                                                  SubscriptionId}}).
+    process_command(#command_unregister_consumer{vhost = VirtualHost,
+                                                 stream = Stream,
+                                                 consumer_name = ConsumerName,
+                                                 connection_pid = ConnectionPid,
+                                                 subscription_id = SubscriptionId}).
 
 -spec activate_consumer(binary(), binary(), binary()) -> ok.
 activate_consumer(VH, Stream, Name) ->
-    process_command({sac,
-                     #command_activate_consumer{vhost =VH,
-                                                stream = Stream,
-                                                consumer_name= Name}}).
+    process_command(#command_activate_consumer{vhost =VH,
+                                               stream = Stream,
+                                               consumer_name= Name}).
 
 -spec connection_reconnected(connection_pid()) -> ok.
 connection_reconnected(Pid) ->
-    process_command({sac,
-                     #command_connection_reconnected{pid = Pid}}).
+    process_command(#command_connection_reconnected{pid = Pid}).
 
 process_command(Cmd) ->
-    case rabbit_stream_coordinator:process_command(Cmd) of
+    case rabbit_stream_coordinator:process_command(wrap_cmd(Cmd)) of
         {ok, Res, _} ->
             Res;
         {error, _} = Err ->
@@ -152,6 +131,10 @@ process_command(Cmd) ->
             Err
     end.
 
+-spec wrap_cmd(command()) -> {sac, command()}.
+wrap_cmd(Cmd) ->
+    {sac, Cmd}.
+
 %% return the current groups for a given virtual host
 -spec consumer_groups(binary(), [atom()]) ->
                          {ok,
@@ -306,8 +289,31 @@ apply(#command_connection_reconnected{pid = Pid},
                           handle_group_connection_reconnected(Pid, St, Eff, G)
                   end, {State0, []}, Groups0),
 
+    {State1, ok, Eff};
+apply(#command_purge_nodes{nodes = Nodes}, State0) ->
+    {State1, Eff} = lists:foldl(fun(N, {S0, Eff0}) ->
+                                        {S1, Eff1} = purge_node(N, S0),
+                                        {S1, Eff1 ++ Eff0}
+                                end, {State0, []}, Nodes),
     {State1, ok, Eff}.
 
+purge_node(Node, #?MODULE{groups = Groups0} = State0) ->
+    PidsGroups =
+        maps:fold(fun(K, #group{consumers = Consumers}, Acc) ->
+                          lists:foldl(fun(#consumer{pid = Pid}, AccIn)
+                                            when node(Pid) =:= Node ->
+                                              PG0 = maps:get(Pid, AccIn, #{}),
+                                              PG1 = PG0#{K => true},
+                                              AccIn#{Pid => PG1};
+                                         (_, AccIn) ->
+                                              AccIn
+                                      end, Acc, Consumers)
+                  end, #{}, Groups0),
+    maps:fold(fun(Pid, Groups, {S0, Eff0}) ->
+                      {S1, Eff1} = handle_connection_down0(Pid, S0, Groups),
+                      {S1, Eff1 ++ Eff0}
+              end, {State0, []}, PidsGroups).
+
 handle_group_connection_reconnected(Pid, #?MODULE{groups = Groups0} = S0,
                                     Eff0, {VH, S, Name} = K) ->
     %% TODO sac: handle forgotten_active case (reconciliate state with current active)
@@ -573,6 +579,7 @@ ensure_monitors(#command_connection_reconnected{pid = Pid},
      Monitors#{Pid => sac},
      [{monitor, process, Pid}, {monitor, node, node(Pid)} | Effects]};
 ensure_monitors(_, #?MODULE{} = State0, Monitors, Effects) ->
+%% TODO sac: ensure the pid-group mapping after purge_nodes?
     {State0, Monitors, Effects}.
 
 -spec handle_connection_down(connection_pid(), state()) ->
@@ -584,11 +591,14 @@ handle_connection_down(Pid,
             {State0, []};
         {Groups, PidsGroups1} ->
             State1 = State0#?MODULE{pids_groups = PidsGroups1},
-            maps:fold(fun(G, _, Acc) ->
-                              handle_group_after_connection_down(Pid, Acc, G)
-                      end, {State1, []}, Groups)
+            handle_connection_down0(Pid, State1, Groups)
     end.
 
+handle_connection_down0(Pid, State, Groups) ->
+    maps:fold(fun(G, _, Acc) ->
+                      handle_group_after_connection_down(Pid, Acc, G)
+              end, {State, []}, Groups).
+
 -spec handle_connection_node_disconnected(connection_pid(), state()) ->
     {state(), ra_machine:effects()}.
 handle_connection_node_disconnected(ConnPid,
@@ -730,6 +740,10 @@ import_state(4, #{<<"groups">> := Groups, <<"pids_groups">> := PidsGroups}) ->
     #?MODULE{groups = map_to_groups(Groups),
              pids_groups = map_to_pids_groups(PidsGroups)}.
 
+- spec make_purge_nodes([node()]) -> command().
+make_purge_nodes(Nodes) ->
+    wrap_cmd(#command_purge_nodes{nodes = Nodes}).
+
 map_to_groups(Groups) when is_map(Groups) ->
     maps:fold(fun(K, V, Acc) ->
                       Acc#{K => map_to_group(V)}
diff --git a/deps/rabbit/src/rabbit_stream_sac_coordinator.hrl b/deps/rabbit/src/rabbit_stream_sac_coordinator.hrl
@@ -64,3 +64,5 @@
          consumer_name :: consumer_name()}).
 -record(command_connection_reconnected,
         {pid :: connection_pid()}).
+-record(command_purge_nodes,
+        {nodes :: [node()]}).
diff --git a/deps/rabbit/test/rabbit_stream_sac_coordinator_SUITE.erl b/deps/rabbit/test/rabbit_stream_sac_coordinator_SUITE.erl
