Re-evaluate stream SAC group after connection down event

The same connection can contain several consumers belonging to a SAC
group (group key = vhost + stream + consumer name). The whole new group
must be re-evaluated to select a new active consumer after the consumers
of the down connection are removed from it.

The previous behavior would not re-evaluate the new group and could
select a consumer from the down connection, letting the group with only
inactive consumers, as the selected active consumer would never receive
the activation message from the stream SAC coordinator.

This commit fixes this problem by removing the consumers of the down
down connection from the affected groups and then performing the
appropriate operations for the groups to keep on consuming (e.g.
notifying an active consumer that it needs to step down).

References #13372

deps/rabbit/src/rabbit_stream_sac_coordinator.erl

@@ -229,7 +229,7 @@ apply(#command_unregister_consumer{vhost = VirtualHost,
                     of
                         {value, Consumer} ->
                             G1 = remove_from_group(Consumer, Group0),
-                            handle_consumer_removal(G1, Consumer, Stream, ConsumerName);
+                            handle_consumer_removal(G1, Stream, ConsumerName, Consumer#consumer.active);
                         false ->
                             {Group0, []}
                     end,
@@ -414,50 +414,44 @@ handle_connection_down(Pid,
             {State0, []};
         {Groups, PidsGroups1} ->
             State1 = State0#?MODULE{pids_groups = PidsGroups1},
-            %% iterate other the groups that this PID affects
+            maps:fold(fun(G, _, Acc) ->
-            maps:fold(fun({VirtualHost, Stream, ConsumerName}, _,
+                              handle_group_after_connection_down(Pid, Acc, G)
-                          {#?MODULE{groups = ConsumerGroups} = S0, Eff0}) ->
+                      end, {State1, []}, Groups)
-                         case lookup_group(VirtualHost,
+    end.
-                                           Stream,
+
-                                           ConsumerName,
+handle_group_after_connection_down(Pid,
-                                           ConsumerGroups)
+                                {#?MODULE{groups = Groups0} = S0, Eff0},
-                         of
+                                {VirtualHost, Stream, ConsumerName}) ->
-                             undefined -> {S0, Eff0};
+    case lookup_group(VirtualHost,
-                             #group{consumers = Consumers} ->
+                      Stream,
-                                 %% iterate over the consumers of the group
+                      ConsumerName,
-                                 %% and unregister the ones from this PID.
+                      Groups0) of
-                                 %% It may not be optimal, computing the new active consumer
+        undefined ->
-                                 %% from the purged group and notifying the remaining consumers
+            {S0, Eff0};
-                                 %% appropriately should avoid unwanted notifications and even rebalancing.
+        #group{consumers = Consumers0} = G0 ->
-                                 lists:foldl(fun (#consumer{pid = P,
+            %% remove the connection consumers from the group state
-                                                            subscription_id =
+            %% keep flags to know what happened
-                                                                SubId},
+            {Consumers1, ActiveRemoved, AnyRemoved} =
-                                                  {StateSub0, EffSub0})
+            lists:foldl(
-                                                     when P == Pid ->
+              fun(#consumer{pid = P, active = S}, {L, ActiveFlag, _}) when P == Pid ->
-                                                     {StateSub1, ok, E} =
+                      {L, S or ActiveFlag, true};
-                                                         ?MODULE:apply(#command_unregister_consumer{vhost
+                 (C, {L, ActiveFlag, AnyFlag}) ->
-                                                                                                        =
+                      {L ++ [C], ActiveFlag, AnyFlag}
-                                                                                                        VirtualHost,
+              end, {[], false, false}, Consumers0),
-                                                                                                    stream
+
-                                                                                                        =
+            case AnyRemoved of
-                                                                                                        Stream,
+                true ->
-                                                                                                    consumer_name
+                    G1 = G0#group{consumers = Consumers1},
-                                                                                                        =
+                    {G2, Effects} = handle_consumer_removal(G1, Stream, ConsumerName, ActiveRemoved),
-                                                                                                        ConsumerName,
+                    Groups1 = update_groups(VirtualHost,
-                                                                                                    connection_pid
+                                            Stream,
-                                                                                                        =
+                                            ConsumerName,
-                                                                                                        Pid,
+                                            G2,
-                                                                                                    subscription_id
+                                            Groups0),
-                                                                                                        =
+                    {S0#?MODULE{groups = Groups1}, Effects ++ Eff0};
-                                                                                                        SubId},
+                false ->
-                                                                       StateSub0),
+                    {S0, Eff0}
-                                                     {StateSub1, EffSub0 ++ E};
+            end
-                                                 (_Consumer, Acc) -> Acc
-                                             end,
-                                             {S0, Eff0}, Consumers)
-                         end
-                      end,
-                      {State1, []}, Groups)
     end.
 
 do_register_consumer(VirtualHost,
@@ -576,9 +570,9 @@ do_register_consumer(VirtualHost,
 handle_consumer_removal(#group{consumers = []} = G, _, _, _) ->
     {G, []};
 handle_consumer_removal(#group{partition_index = -1} = Group0,
-                        Consumer, Stream, ConsumerName) ->
+                        Stream, ConsumerName, ActiveRemoved) ->
-    case Consumer of
+    case ActiveRemoved of
-        #consumer{active = true} ->
+        true ->
             %% this is the active consumer we remove, computing the new one
             Group1 = compute_active_consumer(Group0),
             case lookup_active_consumer(Group1) of
@@ -589,11 +583,11 @@ handle_consumer_removal(#group{partition_index = -1} = Group0,
                     %% no active consumer found in the group, nothing to do
                     {Group1, []}
             end;
-        #consumer{active = false} ->
+        false ->
             %% not the active consumer, nothing to do.
             {Group0, []}
     end;
-handle_consumer_removal(Group0, Consumer, Stream, ConsumerName) ->
+handle_consumer_removal(Group0, Stream, ConsumerName, ActiveRemoved) ->
     case lookup_active_consumer(Group0) of
         {value,
          #consumer{pid = ActPid, subscription_id = ActSubId} =
@@ -612,7 +606,7 @@ handle_consumer_removal(Group0, Consumer, Stream, ConsumerName) ->
                                              Stream, ConsumerName, false, true)]}
             end;
         false ->
-            case Consumer#consumer.active of
+            case ActiveRemoved of
                 true ->
                     %% the active one is going away, picking a new one
                     #consumer{pid = P, subscription_id = SID} =

deps/rabbit/test/rabbit_stream_sac_coordinator_SUITE.erl

@@ -312,29 +312,27 @@ ensure_monitors_test(_) ->
 
     ok.
 
-handle_connection_down_test(_) ->
+handle_connection_down_sac_should_get_activated_test(_) ->
     Stream = <<"stream">>,
     ConsumerName = <<"app">>,
     GroupId = {<<"/">>, Stream, ConsumerName},
     Pid0 = self(),
     Pid1 = spawn(fun() -> ok end),
-    Group =
+    Group = cgroup([consumer(Pid0, 0, true),
-        cgroup([consumer(Pid0, 0, true), consumer(Pid1, 1, false),
+                    consumer(Pid1, 1, false),
-                consumer(Pid0, 2, false)]),
+                    consumer(Pid0, 2, false)]),
-    State0 =
+    State0 = state(#{GroupId => Group},
-        state(#{GroupId => Group},
+                   #{Pid0 => maps:from_list([{GroupId, true}]),
-              #{Pid0 => maps:from_list([{GroupId, true}]),
+                     Pid1 => maps:from_list([{GroupId, true}])}),
-                Pid1 => maps:from_list([{GroupId, true}])}),
 
     {#?STATE{pids_groups = PidsGroups1, groups = Groups1} = State1,
      Effects1} =
-        rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State0),
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State0),
     assertSize(1, PidsGroups1),
     assertSize(1, maps:get(Pid1, PidsGroups1)),
     assertSendMessageEffect(Pid1, 1, Stream, ConsumerName, true, Effects1),
-    ?assertEqual(#{GroupId => cgroup([consumer(Pid1, 1, true)])},
+    assertHasGroup(GroupId, cgroup([consumer(Pid1, 1, true)]), Groups1),
-                 Groups1),
+    {#?STATE{pids_groups = PidsGroups2, groups = Groups2},
-    {#?STATE{pids_groups = PidsGroups2, groups = Groups2} = _State2,
      Effects2} =
         rabbit_stream_sac_coordinator:handle_connection_down(Pid1, State1),
     assertEmpty(PidsGroups2),
@@ -343,6 +341,168 @@ handle_connection_down_test(_) ->
 
     ok.
 
+handle_connection_down_sac_active_does_not_change_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    Group = cgroup([consumer(Pid1, 0, true),
+                    consumer(Pid0, 1, false),
+                    consumer(Pid0, 2, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]),
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid1, PidsGroups)),
+    assertEmpty(Effects),
+    assertHasGroup(GroupId, cgroup([consumer(Pid1, 0, true)]), Groups),
+    ok.
+
+handle_connection_down_sac_no_more_consumers_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Group = cgroup([consumer(Pid0, 0, true),
+                    consumer(Pid0, 1, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+    assertEmpty(PidsGroups),
+    assertEmpty(Groups),
+    assertEmpty(Effects),
+    ok.
+
+handle_connection_down_sac_no_consumers_in_down_connection_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    Group = cgroup([consumer(Pid1, 0, true),
+                    consumer(Pid1, 1, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]), %% should not be there
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid1, PidsGroups)),
+    assertEmpty(Effects),
+    assertHasGroup(GroupId, cgroup([consumer(Pid1, 0, true), consumer(Pid1, 1, false)]),
+                   Groups),
+    ok.
+
+handle_connection_down_super_stream_active_stays_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    Group = cgroup(1, [consumer(Pid0, 0, false),
+                       consumer(Pid0, 1, true),
+                       consumer(Pid1, 2, false),
+                       consumer(Pid1, 3, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]),
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid1, State),
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid0, PidsGroups)),
+    assertEmpty(Effects),
+    assertHasGroup(GroupId, cgroup(1, [consumer(Pid0, 0, false), consumer(Pid0, 1, true)]),
+                   Groups),
+    ok.
+
+handle_connection_down_super_stream_active_changes_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    Group = cgroup(1, [consumer(Pid0, 0, false),
+                       consumer(Pid1, 1, true),
+                       consumer(Pid0, 2, false),
+                       consumer(Pid1, 3, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]),
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid1, PidsGroups)),
+    assertSendMessageSteppingDownEffect(Pid1, 1, Stream, ConsumerName, Effects),
+    assertHasGroup(GroupId, cgroup(1, [consumer(Pid1, 1, false), consumer(Pid1, 3, false)]),
+                   Groups),
+    ok.
+
+handle_connection_down_super_stream_activate_in_remaining_connection_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    Group = cgroup(1, [consumer(Pid0, 0, false),
+                       consumer(Pid0, 1, true),
+                       consumer(Pid1, 2, false),
+                       consumer(Pid1, 3, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]),
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid1, PidsGroups)),
+    assertSendMessageEffect(Pid1, 3, Stream, ConsumerName, true, Effects),
+    assertHasGroup(GroupId, cgroup(1, [consumer(Pid1, 2, false), consumer(Pid1, 3, true)]),
+                   Groups),
+    ok.
+
+handle_connection_down_super_stream_no_active_removed_or_present_test(_) ->
+    Stream = <<"stream">>,
+    ConsumerName = <<"app">>,
+    GroupId = {<<"/">>, Stream, ConsumerName},
+    Pid0 = self(),
+    Pid1 = spawn(fun() -> ok end),
+    %% this is a weird case that should not happen in the wild,
+    %% we test the logic in the code nevertheless.
+    %% No active consumer in the group
+    Group = cgroup(1, [consumer(Pid0, 0, false),
+                       consumer(Pid0, 1, false),
+                       consumer(Pid1, 2, false),
+                       consumer(Pid1, 3, false)]),
+    State = state(#{GroupId => Group},
+                  #{Pid0 => maps:from_list([{GroupId, true}]),
+                    Pid1 => maps:from_list([{GroupId, true}])}),
+
+    {#?STATE{pids_groups = PidsGroups, groups = Groups},
+     Effects} =
+    rabbit_stream_sac_coordinator:handle_connection_down(Pid0, State),
+    assertSize(1, PidsGroups),
+    assertSize(1, maps:get(Pid1, PidsGroups)),
+    assertEmpty(Effects),
+    assertHasGroup(GroupId, cgroup(1, [consumer(Pid1, 2, false), consumer(Pid1, 3, false)]),
+                   Groups),
+    ok.
+
 assertSize(Expected, []) ->
     ?assertEqual(Expected, 0);
 assertSize(Expected, Map) when is_map(Map) ->
@@ -353,6 +513,9 @@ assertSize(Expected, List) when is_list(List) ->
 assertEmpty(Data) ->
     assertSize(0, Data).
 
+assertHasGroup(GroupId, Group, Groups) ->
+    ?assertEqual(#{GroupId => Group}, Groups).
+
 consumer(Pid, SubId, Active) ->
     #consumer{pid = Pid,
               subscription_id = SubId,

deps/rabbitmq_stream/src/rabbit_stream_reader.erl

@@ -598,26 +598,16 @@ augment_infos_with_user_provided_connection_name(Infos,
     end.
 
 close(Transport,
-      #stream_connection{socket = S, virtual_host = VirtualHost,
+      #stream_connection{socket = S},
-                         outstanding_requests = Requests},
       #stream_connection_state{consumers = Consumers}) ->
     [begin
-         %% we discard the result (updated requests) because they are no longer used
-         _ = maybe_unregister_consumer(VirtualHost, Consumer,
-                                       single_active_consumer(Properties),
-                                       Requests),
          case Log of
              undefined ->
                  ok; %% segment may not be defined on subscription (single active consumer)
              L ->
                  osiris_log:close(L)
          end
-     end
+     end || #consumer{log = Log} <- maps:values(Consumers)],
-     || #consumer{log = Log,
-                  configuration =
-                      #consumer_configuration{properties = Properties}} =
-            Consumer
-            <- maps:values(Consumers)],
     Transport:shutdown(S, write),
     Transport:close(S).