Implement node purge in stream SAC coordinator

Author: acogoluegnes

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) ->

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)}

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()]}).

deps/rabbit/test/rabbit_stream_sac_coordinator_SUITE.erl

@@ -646,7 +646,7 @@ connection_reconnected_simple_disconnected_becomes_connected_test(_) ->
     Groups0 = #{GId => Group},
     State0 = state(Groups0),
 
-    Cmd = connection_reconnection_command(Pid0),
+    Cmd = connection_reconnected_command(Pid0),
     {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
 
     assertHasGroup(GId, cgroup([consumer(Pid0, 0, {connected, active}),
@@ -670,7 +670,7 @@ connection_reconnected_simple_active_should_be_first_test(_) ->
     Groups0 = #{GId => Group},
     State0 = state(Groups0),
 
-    Cmd = connection_reconnection_command(Pid0),
+    Cmd = connection_reconnected_command(Pid0),
     {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
 
     assertHasGroup(GId, cgroup([consumer(Pid1, 1, {connected, active}),
@@ -692,7 +692,7 @@ connection_reconnected_super_disconnected_becomes_connected_test(_) ->
     Groups0 = #{GId => Group},
     State0 = state(Groups0),
 
-    Cmd = connection_reconnection_command(Pid0),
+    Cmd = connection_reconnected_command(Pid0),
     {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
 
     assertHasGroup(GId, cgroup(1, [consumer(Pid0, 0, {connected, waiting}),
@@ -724,7 +724,7 @@ forget_connection_simple_disconnected_becomes_forgotten_test(_) ->
     assertSendMessageEffect(Pid1, 1, stream(), name(), true, Eff),
     ok.
 
-forget_connection_super_disconnected_becomes_forgotten_test(_) ->
+forget_connection_super_stream_disconnected_becomes_forgotten_test(_) ->
     Pid0 = spawn(fun() -> ok end),
     Pid1 = spawn(fun() -> ok end),
     Pid2 = spawn(fun() -> ok end),
@@ -936,8 +936,153 @@ handle_connection_node_disconnected_super_stream_disconn_active_block_rebalancin
     assertNodeDisconnectedTimerEffect(Pid0, Eff),
     ok.
 
+connection_reconnected_simple_disconn_active_block_rebalancing_test(_) ->
+    Pid0 = spawn(fun() -> ok end),
+    Pid1 = spawn(fun() -> ok end),
+    Pid2 = spawn(fun() -> ok end),
+    GId = group_id(),
+    Group = cgroup([consumer(Pid0, 0, {disconnected, waiting}),
+                    consumer(Pid1, 0, {disconnected, active}),
+                    consumer(Pid2, 0, {connected, waiting})]),
+
+    Groups0 = #{GId => Group},
+    State0 = state(Groups0),
+    Cmd = connection_reconnected_command(Pid0),
+    {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
+
+    assertHasGroup(GId, cgroup([consumer(Pid1, 0, {disconnected, active}),
+                                consumer(Pid0, 0, {connected, waiting}),
+                                consumer(Pid2, 0, {connected, waiting})]),
+                   Groups1),
+    assertEmpty(Eff),
+    ok.
+
+connection_reconnected_super_stream_disconn_active_block_rebalancing_test(_) ->
+    Pid0 = spawn(fun() -> ok end),
+    Pid1 = spawn(fun() -> ok end),
+    Pid2 = spawn(fun() -> ok end),
+    GId = group_id(),
+    Group = cgroup(1, [consumer(Pid0, 0, {disconnected, active}),
+                       consumer(Pid1, 0, {disconnected, waiting}),
+                       consumer(Pid2, 0, {connected, waiting})]),
+
+    Groups0 = #{GId => Group},
+    State0 = state(Groups0),
+    Cmd = connection_reconnected_command(Pid1),
+    {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
+
+    assertHasGroup(GId, cgroup(1, [consumer(Pid0, 0, {disconnected, active}),
+                                   consumer(Pid1, 0, {connected, waiting}),
+                                   consumer(Pid2, 0, {connected, waiting})]),
+                   Groups1),
+    assertEmpty(Eff),
+    ok.
+
+forget_connection_simple_disconn_active_block_rebalancing_test(_) ->
+    Pid0 = spawn(fun() -> ok end),
+    Pid1 = spawn(fun() -> ok end),
+    Pid2 = spawn(fun() -> ok end),
+    GId = group_id(),
+    Group = cgroup([consumer(Pid0, {disconnected, waiting}),
+                    consumer(Pid1, {connected, waiting}),
+                    consumer(Pid2, {disconnected, active})]),
+
+    Groups0 = #{GId => Group},
+    State0 = state(Groups0),
+
+    {#?STATE{groups = Groups1}, Eff} = ?MOD:forget_connection(Pid0, State0),
+
+    assertHasGroup(GId, cgroup([consumer(Pid2, {disconnected, active}),
+                                consumer(Pid0, {forgotten, waiting}),
+                                consumer(Pid1, {connected, waiting})]),
+                   Groups1),
+    assertEmpty(Eff),
+    ok.
+
+forget_connection_super_stream_disconn_active_block_rebalancing_test(_) ->
+    Pid0 = spawn(fun() -> ok end),
+    Pid1 = spawn(fun() -> ok end),
+    Pid2 = spawn(fun() -> ok end),
+    GId = group_id(),
+    Group = cgroup(1, [consumer(Pid0, {disconnected, waiting}),
+                       consumer(Pid1, {connected, waiting}),
+                       consumer(Pid2, {disconnected, active})]),
+
+    Groups0 = #{GId => Group},
+    State0 = state(Groups0),
+
+    {#?STATE{groups = Groups1}, Eff} = ?MOD:forget_connection(Pid0, State0),
+
+    assertHasGroup(GId, cgroup(1, [consumer(Pid0, {forgotten, waiting}),
+                                   consumer(Pid1, {connected, waiting}),
+                                   consumer(Pid2, {disconnected, active})]),
+                   Groups1),
+    assertEmpty(Eff),
+    ok.
+
+purge_nodes_test(_) ->
+    N0 = node(),
+    {ok, N1Pid, N1} = peer:start(#{
+        name => ?FUNCTION_NAME,
+        connection => standard_io,
+        shutdown => close
+    }),
+
+    N0P0 = spawn(N0, fun() -> ok end),
+    N0P1 = spawn(N0, fun() -> ok end),
+    N0P2 = spawn(N0, fun() -> ok end),
+    N1P0 = spawn(N1, fun() -> ok end),
+    N1P1 = spawn(N1, fun() -> ok end),
+    N1P2 = spawn(N1, fun() -> ok end),
+
+    S0 = <<"s0">>,
+    S1 = <<"s1">>,
+    S2 = <<"s2">>,
+
+    GId0 = group_id(S0),
+    GId1 = group_id(S1),
+    GId2 = group_id(S2),
+
+    Group0 = cgroup([consumer(N1P0, {disconnected, active}),
+                     consumer(N0P1, {connected, waiting}),
+                     consumer(N0P2, {connected, waiting})]),
+
+    Group1 = cgroup(1, [consumer(N1P1, {disconnected, waiting}),
+                        consumer(N1P2, {disconnected, active}),
+                        consumer(N0P0, {connected, waiting})]),
+
+    Group2 = cgroup([consumer(N0P0, {connected, active}),
+                     consumer(N0P1, {connected, waiting}),
+                     consumer(N0P2, {connected, waiting})]),
+
+
+    State0 = state(#{GId0 => Group0, GId1 => Group1, GId2 => Group2}),
+    Cmd = purge_nodes_command([N1]),
+    {#?STATE{groups = Groups1}, ok, Eff} = ?MOD:apply(Cmd, State0),
+
+    assertSize(3, Groups1),
+    assertHasGroup(GId0, cgroup([consumer(N0P1, {connected, active}),
+                                 consumer(N0P2, {connected, waiting})]),
+                   Groups1),
+    assertHasGroup(GId1, cgroup(1, [consumer(N0P0, {connected, active})]),
+                   Groups1),
+    assertHasGroup(GId2, cgroup([consumer(N0P0, {connected, active}),
+                                 consumer(N0P1, {connected, waiting}),
+                                 consumer(N0P2, {connected, waiting})]),
+                   Groups1),
+
+    assertSize(2, Eff),
+    assertContainsSendMessageEffect(N0P1, S0, true, Eff),
+    assertContainsSendMessageEffect(N0P0, S1, true, Eff),
+
+    _ = peer:stop(N1Pid),
+    ok.
+
 group_id() ->
-    {<<"/">>, stream(), name()}.
+    group_id(stream()).
+
+group_id(S) ->
+    {<<"/">>, S, name()}.
 
 stream() ->
     <<"sO">>.
@@ -964,6 +1109,9 @@ assertHasGroup(GroupId, Group, Groups) ->
     G = maps:get(GroupId, Groups),
     ?assertEqual(Group, G).
 
+consumer(Pid, Status) ->
+    consumer(Pid, 0, Status).
+
 consumer(Pid, SubId, {Connectivity, Status}) ->
     #consumer{pid = Pid,
               subscription_id = SubId,
@@ -1015,9 +1163,32 @@ activate_consumer_command(Stream, ConsumerName) ->
                                stream = Stream,
                                consumer_name = ConsumerName}.
 
-connection_reconnection_command(Pid) ->
+connection_reconnected_command(Pid) ->
     #command_connection_reconnected{pid = Pid}.
 
+purge_nodes_command(Nodes) ->
+    #command_purge_nodes{nodes = Nodes}.
+
+
+assertContainsSendMessageEffect(Pid, Stream, Active, Effects) ->
+    assertContainsSendMessageEffect(Pid, 0, Stream, name(), Active, Effects).
+
+assertContainsSendMessageEffect(Pid, SubId, Stream, ConsumerName, Active,
+                                Effects) ->
+    Contains = lists:any(fun(Eff) ->
+                                 Eff =:= {mod_call,
+                                          rabbit_stream_sac_coordinator,
+                                          send_message,
+                                          [Pid,
+                                           {sac,
+                                            #{subscription_id => SubId,
+                                              stream => Stream,
+                                              consumer_name => ConsumerName,
+                                              active => Active}}]}
+              end, Effects),
+    ?assert(Contains).
+
+
 assertSendMessageEffect(Pid, SubId, Stream, ConsumerName, Active, [Effect]) ->
     ?assertEqual({mod_call,
                   rabbit_stream_sac_coordinator,