Merge pull request #4496 from esl/debug-stream-management

DenysGonchar · web-flow · commit a58666850c52 · 2025-03-10T16:28:30.000+01:00
Deduplicate stream management buffer (solution 2+3)
diff --git a/big_tests/tests/sm_SUITE.erl b/big_tests/tests/sm_SUITE.erl
@@ -125,7 +125,10 @@ parallel_cases() ->
 parallel_large_buffer_cases() ->
     [resend_unacked_from_stopped_sessions,
      resend_unacked_from_terminated_sessions,
-     resend_unacked_from_replaced_sessions].
+     resend_unacked_from_replaced_sessions,
+     relay_unacked_from_stopped_sessions,
+     relay_unacked_from_terminated_sessions,
+     relay_unacked_from_replaced_sessions].
 
 parallel_manual_ack_freq_1_cases() ->
     [client_acks_more_than_sent,
@@ -591,24 +594,30 @@ resend_more_offline_messages_than_buffer_size(Config) ->
 
 %% Test cases for duplicate buffer
 
+-define(USER_NUM, 4).
+
 resend_unacked_from_stopped_sessions(Config) ->
     Texts = [<<"msg-1">>],
-    {Bob, UserSpecs, Users} = connect_initial_users(Texts, Config),
+    {Bob, UserSpecs, Users} = connect_initial_users(Config),
 
     %% Bob sends messages to User's bare jid
     UserJid = escalus_users:get_jid(Config, hd(UserSpecs)),
     sm_helper:send_messages(Bob, UserJid, Texts),
     C2SPids = lists:map(fun mongoose_helper:get_session_pid/1, Users),
     [sm_helper:wait_for_c2s_unacked_count(Pid, length(Texts)) || Pid <- C2SPids],
 
-    %% Each User session checks messages and stops,
-    %% causing buffer rerouting and message duplication
-    Funs = [fun() -> User end || User <- Users],
-    reconnect_and_receive_messages(Funs, Texts). % Reconnection is skipped in this test case
+    %% Each User's session checks messages and stops, rerouting unacked messages to online sessions.
+    %% This is why there is one extra copy of each message received in each subsequent iteration.
+    lists:foreach(fun({N, User}) ->
+                          DuplicatedTexts = lists:append(lists:duplicate(N, Texts)),
+                          receive_unacked_messages(User, DuplicatedTexts),
+                          escalus_connection:stop(User)
+                  end, lists:enumerate(Users)),
+    escalus_connection:stop(Bob).
 
 resend_unacked_from_terminated_sessions(Config) ->
     Texts = [<<"msg-1">>],
-    {Bob, UserSpecs, Users} = connect_initial_users(Texts, Config),
+    {Bob, UserSpecs, Users} = connect_initial_users(Config),
 
     %% User disconnects all sessions abruptly
     lists:foreach(fun escalus_connection:kill/1, Users),
@@ -621,13 +630,17 @@ resend_unacked_from_terminated_sessions(Config) ->
     [sm_helper:wait_for_c2s_unacked_count(Pid, length(Texts)) || Pid <- C2SPids],
 
     %% User replaces each terminated session with a new one,
-    %% causing buffer rerouting and message duplication
-    Funs = [fun() -> connect_spec(Spec, session) end || Spec <- UserSpecs],
-    reconnect_and_receive_messages(Funs, Texts).
+    %% rerouting unacked messages to the new session
+    lists:foreach(fun(UserSpec) ->
+                          NewUser = connect_spec(UserSpec, session),
+                          receive_unacked_messages(NewUser, Texts),
+                          escalus_connection:stop(NewUser)
+                  end, UserSpecs),
+    escalus_connection:stop(Bob).
 
 resend_unacked_from_replaced_sessions(Config) ->
     Texts = [<<"msg-1">>],
-    {Bob, UserSpecs, Users} = connect_initial_users(Texts, Config),
+    {Bob, UserSpecs, Users} = connect_initial_users(Config),
 
     %% Bob sends messages to User's bare jid
     UserJid = escalus_users:get_jid(Config, hd(UserSpecs)),
@@ -636,32 +649,91 @@ resend_unacked_from_replaced_sessions(Config) ->
     [sm_helper:wait_for_c2s_unacked_count(Pid, length(Texts)) || Pid <- C2SPids],
 
     %% User replaces each online session with a new one,
-    %% causing buffer rerouting and message duplication
-    Funs = [fun() -> connect_spec(Spec, session) end || Spec <- UserSpecs],
-    reconnect_and_receive_messages(Funs, Texts).
+    %% rerouting unacked messages to the new session
+    lists:foreach(fun(UserSpec) ->
+                          NewUser = connect_spec(UserSpec, session),
+                          receive_unacked_messages(NewUser, Texts),
+                          escalus_connection:stop(NewUser)
+                  end, UserSpecs),
+    escalus_connection:stop(Bob).
 
-connect_initial_users(Texts, Config) ->
-    Resources = [<<"res-", (integer_to_binary(I))/binary>> || I <- lists:seq(1, 4)],
+connect_initial_users(Config) ->
+    Resources = [<<"res-", (integer_to_binary(I))/binary>> || I <- lists:seq(1, ?USER_NUM)],
     Bob = connect_fresh(Config, bob, session),
     BasicUserSpec = escalus_fresh:create_fresh_user(Config, ?config(user, Config)),
     UserSpecs = [[{resource, Res} | BasicUserSpec] || Res <- Resources],
     Users = [connect_spec(Spec, sm_after_session) || Spec <- UserSpecs],
     {Bob, UserSpecs, Users}.
 
-%% Reconnect (optionally), receive messages and disconnect each resource cleanly, in sequence
-%% There is an issue causing each subsequent session
-%% to receive twice as many messages as the previous one, i.e. 1, 2, 4, 8, ...
-%% See https://github.com/esl/MongooseIM/pull/4498 for a more detailed description with a diagram
-reconnect_and_receive_messages([UserF | Rest], Texts) ->
-    NewUser = UserF(),
-    sm_helper:wait_for_messages(NewUser, Texts),
-    timer:sleep(100), % wait a short time to ensure no extra messages arrive
-    escalus_assert:has_no_stanzas(NewUser),
-    escalus_connection:stop(NewUser),
-    %% Expect duplicated buffer in the next session
-    reconnect_and_receive_messages(Rest, Texts ++ Texts);
-reconnect_and_receive_messages([], _Texts) ->
-    ok.
+relay_unacked_from_stopped_sessions(Config) ->
+    Texts = [<<"msg-1">>],
+    Resources = [<<"res-", (integer_to_binary(I))/binary>> || I <- lists:seq(1, ?USER_NUM)],
+    Bob = connect_fresh(Config, bob, session),
+    BasicUserSpec = escalus_fresh:create_fresh_user(Config, ?config(user, Config)),
+    UserSpecs = [[{resource, Res} | BasicUserSpec] || Res <- Resources],
+    FirstUser = connect_spec(hd(UserSpecs), sm_after_session),
+
+    %% Bob sends messages to User's bare jid
+    UserJid = escalus_users:get_jid(Config, hd(UserSpecs)),
+    sm_helper:send_messages(Bob, UserJid, Texts),
+
+    %% Each User's session resends unacked messages to the next one
+    RelayF = fun(NextSpec, CurUser) ->
+                     receive_unacked_messages(CurUser, Texts),
+                     NewUser = connect_spec(NextSpec, sm_before_session),
+                     escalus_connection:stop(CurUser),
+                     NewUser
+             end,
+    LastUser = lists:foldl(RelayF, FirstUser, tl(UserSpecs)),
+    escalus_connection:stop(LastUser),
+    escalus_connection:stop(Bob).
+
+relay_unacked_from_terminated_sessions(Config) ->
+    Texts = [<<"msg-1">>],
+    Bob = connect_fresh(Config, bob, session),
+    UserSpec = escalus_fresh:create_fresh_user(Config, ?config(user, Config)),
+    FirstUser = connect_spec(UserSpec, sm_after_session),
+
+    %% Bob sends messages to User's bare jid
+    UserJid = escalus_users:get_jid(Config, UserSpec),
+    sm_helper:send_messages(Bob, UserJid, Texts),
+
+    %% Each User's session resends unacked messages to the next one
+    RelayF = fun(NextSpec, CurUser) ->
+                     receive_unacked_messages(CurUser, Texts),
+                     escalus_connection:kill(CurUser),
+                     C2SPid = mongoose_helper:get_session_pid(CurUser),
+                     sm_helper:wait_until_resume_session(C2SPid),
+                     connect_spec(NextSpec, sm_before_session)
+             end,
+    LastUser = lists:foldl(RelayF, FirstUser, lists:duplicate(?USER_NUM - 1, UserSpec)),
+    escalus_connection:stop(LastUser),
+    escalus_connection:stop(Bob).
+
+relay_unacked_from_replaced_sessions(Config) ->
+    Texts = [<<"msg-1">>],
+    Bob = connect_fresh(Config, bob, session),
+    UserSpec = escalus_fresh:create_fresh_user(Config, ?config(user, Config)),
+    FirstUser = connect_spec(UserSpec, sm_after_session),
+
+    %% Bob sends messages to User's bare jid
+    UserJid = escalus_users:get_jid(Config, UserSpec),
+    sm_helper:send_messages(Bob, UserJid, Texts),
+
+    %% Each User's session resends unacked messages to the next one
+    RelayF = fun(NextSpec, CurUser) ->
+                     receive_unacked_messages(CurUser, Texts),
+                     connect_spec(NextSpec, sm_before_session)
+             end,
+    LastUser = lists:foldl(RelayF, FirstUser, lists:duplicate(?USER_NUM - 1, UserSpec)),
+    escalus_connection:stop(LastUser),
+    escalus_connection:stop(Bob).
+
+%% Receive expected messages and wait a bit to ensure no extra messages arrive
+receive_unacked_messages(User, Texts) ->
+    sm_helper:wait_for_messages(User, Texts),
+    timer:sleep(100),
+    escalus_assert:has_no_stanzas(User).
 
 resend_unacked_on_reconnection(Config) ->
     Texts = three_texts(),
diff --git a/src/c2s/mongoose_c2s.erl b/src/c2s/mongoose_c2s.erl
@@ -662,6 +662,12 @@ maybe_retry_state(StateData = #c2s_data{listener_opts = LOpts}, C2SState) ->
     end.
 
 -spec handle_cast(data(), state(), term()) -> fsm_res().
+handle_cast(StateData, _C2SState, {exit, {replaced, _} = Reason}) ->
+    ReasonText = <<"Replaced by new connection">>,
+    StreamConflict = mongoose_xmpp_errors:stream_conflict(StateData#c2s_data.lang, ReasonText),
+    send_element_from_server_jid(StateData, StreamConflict),
+    send_trailer(StateData),
+    {stop, {shutdown, Reason}};
 handle_cast(StateData, _C2SState, {exit, Reason}) when is_binary(Reason) ->
     StreamConflict = mongoose_xmpp_errors:stream_conflict(StateData#c2s_data.lang, Reason),
     send_element_from_server_jid(StateData, StreamConflict),
@@ -1131,7 +1137,7 @@ start_link(Params, ProcOpts) ->
 stop(Pid, Reason) ->
     gen_statem:cast(Pid, {stop, Reason}).
 
--spec exit(pid(), binary() | atom()) -> ok.
+-spec exit(pid(), binary() | atom() | {replaced, pid()}) -> ok.
 exit(Pid, Reason) ->
     gen_statem:cast(Pid, {exit, Reason}).
 
diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl
@@ -112,7 +112,7 @@
 -type info() :: #{info_key() => any()}.
 
 -type backend() :: ejabberd_sm_mnesia | ejabberd_sm_redis | ejabberd_sm_cets.
--type close_reason() :: resumed | normal | replaced.
+-type close_reason() :: resumed | normal | {replaced, pid()}.
 -type info_key() :: atom().
 
 -export_type([session/0,
@@ -193,7 +193,7 @@ make_new_sid() ->
       ReplacedPids :: [pid()].
 open_session(HostType, SID, JID, Priority, Info) ->
     set_session(SID, JID, Priority, Info),
-    ReplacedPIDs = check_for_sessions_to_replace(HostType, JID),
+    ReplacedPIDs = check_for_sessions_to_replace(HostType, SID, JID),
     mongoose_instrument:execute(sm_session, #{host_type => HostType},
                                 #{jid => JID, logins => 1, count => 1}),
     mongoose_hooks:sm_register_connection(HostType, SID, JID, Info),
@@ -871,18 +871,19 @@ is_offline(#jid{luser = LUser, lserver = LServer}) ->
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 
 %% @doc On new session, check if some existing connections need to be replace
--spec check_for_sessions_to_replace(HostType, JID) -> ReplacedPids when
+-spec check_for_sessions_to_replace(HostType, SID, JID) -> ReplacedPids when
       HostType :: mongooseim:host_type(),
+      SID :: sid(),
       JID :: jid:jid(),
       ReplacedPids :: [pid()].
-check_for_sessions_to_replace(HostType, JID) ->
+check_for_sessions_to_replace(HostType, {_, NewPid}, JID) ->
     #jid{luser = LUser, lserver = LServer, lresource = LResource} = JID,
     Sessions = ejabberd_sm_backend:get_sessions(LUser, LServer),
     %% TODO: Depending on how this is executed, there could be an unneeded
     %% replacement for max_sessions. We need to check this at some point.
     ReplacedRedundantSessions = check_existing_resources(LResource, Sessions),
     AllReplacedSessionPids = check_max_sessions(HostType, LUser, LServer, ReplacedRedundantSessions, Sessions),
-    [mongoose_c2s:exit(Pid, <<"Replaced by new connection">>) || Pid <- AllReplacedSessionPids],
+    [mongoose_c2s:exit(Pid, {replaced, NewPid}) || Pid <- AllReplacedSessionPids],
     AllReplacedSessionPids.
 
 -spec check_existing_resources(LResource, Sessions) ->
diff --git a/src/mod_presence.erl b/src/mod_presence.erl
@@ -560,7 +560,7 @@ close_session_status(normal) ->
     <<>>;
 close_session_status({shutdown, retries}) ->
     <<"Too many attempts">>;
-close_session_status({shutdown, replaced}) ->
+close_session_status({shutdown, {replaced, _Pid}}) ->
     <<"Replaced by new connection">>;
 close_session_status({shutdown, Reason}) when is_atom(Reason) ->
     <<"Shutdown by reason: ", (atom_to_binary(Reason))/binary>>;
diff --git a/src/stream_management/mod_stream_management.erl b/src/stream_management/mod_stream_management.erl
@@ -43,7 +43,6 @@
 -include("mongoose_config_spec.hrl").
 -define(STREAM_MGMT_H_MAX, (1 bsl 32 - 1)).
 -define(CONSTRAINT_CHECK_TIMEOUT, 5000).  %% 5 seconds
--define(IS_STREAM_MGMT_STOP(R), R =:= {shutdown, ?MODULE}; R =:= {shutdown, resumed}).
 -define(IS_ALLOWED_STATE(S), S =:= wait_for_session_establishment; S =:= session_established).
 
 -record(sm_state, {
@@ -229,8 +228,27 @@ xmpp_presend_element(Acc, #{c2s_data := StateData, c2s_state := C2SState}, _Extr
         {#sm_state{buffer_max = no_buffer} = SmState, _} ->
             maybe_send_ack_request(Acc, C2SState, SmState);
         {SmState, _} ->
-            Jid = mongoose_c2s:get_jid(StateData),
-            handle_buffer_and_ack(Acc, C2SState, Jid, SmState)
+            case prevent_buffer_duplication(Acc, SmState) of
+                duplicate ->
+                    {stop, Acc};
+                Acc1 ->
+                    Jid = mongoose_c2s:get_jid(StateData),
+                    handle_buffer_and_ack(Acc1, C2SState, Jid, SmState)
+            end
+    end.
+
+-spec prevent_buffer_duplication(mongoose_acc:t(), sm_state()) -> duplicate | mongoose_acc:t().
+prevent_buffer_duplication(Acc, #sm_state{buffer = Buffer}) ->
+    case mongoose_acc:get(?MODULE, buffer_ref, undefined, Acc) of
+        undefined ->
+            mongoose_acc:set_permanent(?MODULE, buffer_ref, make_ref(), Acc);
+        Ref ->
+            case lists:any(fun(BufAcc) ->
+                                   mongoose_acc:get(?MODULE, buffer_ref, undefined, BufAcc) =:= Ref
+                           end, Buffer) of
+                true -> duplicate;
+                false -> Acc
+            end
     end.
 
 -spec handle_buffer_and_ack(mongoose_acc:t(), c2s_state(), jid:jid(), sm_state()) ->
@@ -362,44 +380,45 @@ notify_unacknowledged_msg(Acc, Jid) ->
 
 -spec reroute_unacked_messages(mongoose_acc:t(), mongoose_c2s_hooks:params(), gen_hook:extra()) ->
     mongoose_c2s_hooks:result().
-reroute_unacked_messages(Acc, #{c2s_data := StateData, reason := Reason}, #{host_type := HostType}) ->
+reroute_unacked_messages(Acc, #{c2s_state := C2SState, c2s_data := StateData, reason := Reason}, #{host_type := HostType}) ->
     MaybeSmState = get_mod_state(StateData),
-    maybe_handle_stream_mgmt_reroute(Acc, StateData, HostType, Reason, MaybeSmState).
+    maybe_handle_stream_mgmt_reroute(Acc, C2SState, StateData, HostType, Reason, MaybeSmState).
 
 -spec user_terminate(mongoose_acc:t(), mongoose_c2s_hooks:params(), gen_hook:extra()) ->
     mongoose_c2s_hooks:result().
-user_terminate(Acc, #{reason := Reason}, _Extra) when ?IS_STREAM_MGMT_STOP(Reason) ->
+user_terminate(Acc, #{c2s_state := ?EXT_C2S_STATE(resume_session)}, _Extra) ->
     {stop, Acc}; %% We stop here because this termination was triggered internally
 user_terminate(Acc, _Params, _Extra) ->
     {ok, Acc}.
 
 -spec maybe_handle_stream_mgmt_reroute(
-        mongoose_acc:t(), mongoose_c2s:data(), mongooseim:host_type(), term(), maybe_sm_state()) ->
+        mongoose_acc:t(), mongoose_c2s:state(), mongoose_c2s:data(), mongooseim:host_type(), term(), maybe_sm_state()) ->
     mongoose_c2s_hooks:result().
-maybe_handle_stream_mgmt_reroute(Acc, StateData, HostType, Reason, #sm_state{counter_in = H} = SmState)
-  when ?IS_STREAM_MGMT_STOP(Reason) ->
+maybe_handle_stream_mgmt_reroute(Acc, ?EXT_C2S_STATE(resume_session), StateData, HostType, Reason, #sm_state{counter_in = H} = SmState) ->
     Sid = mongoose_c2s:get_sid(StateData),
     do_remove_smid(HostType, Sid, H),
-    NewSmState = handle_user_terminate(SmState, StateData, HostType),
+    NewSmState = handle_user_terminate(SmState, StateData, HostType, Reason),
     {ok, mongoose_c2s_acc:to_acc(Acc, state_mod, {?MODULE, NewSmState})};
-maybe_handle_stream_mgmt_reroute(Acc, StateData, HostType, _Reason, #sm_state{} = SmState) ->
-    NewSmState = handle_user_terminate(SmState, StateData, HostType),
+maybe_handle_stream_mgmt_reroute(Acc, _C2SState, StateData, HostType, Reason, #sm_state{} = SmState) ->
+    NewSmState = handle_user_terminate(SmState, StateData, HostType, Reason),
     {ok, mongoose_c2s_acc:to_acc(Acc, state_mod, {?MODULE, NewSmState})};
-maybe_handle_stream_mgmt_reroute(Acc, _StateData, _HostType, _Reason, {error, not_found}) ->
+maybe_handle_stream_mgmt_reroute(Acc, _C2SState,_StateData, _HostType, _Reason, {error, not_found}) ->
     {ok, Acc}.
 
--spec handle_user_terminate(sm_state(), mongoose_c2s:data(), mongooseim:host_type()) -> sm_state().
-handle_user_terminate(#sm_state{counter_in = H} = SmState, StateData, HostType) ->
+-spec handle_user_terminate(sm_state(), mongoose_c2s:data(), mongooseim:host_type(), term()) -> sm_state().
+handle_user_terminate(#sm_state{counter_in = H} = SmState, StateData, HostType, Reason) ->
     Sid = mongoose_c2s:get_sid(StateData),
     do_remove_smid(HostType, Sid, H),
     FromServer = mongoose_c2s:get_lserver(StateData),
     NewState = add_delay_elements_to_buffer(SmState, FromServer),
-    reroute_buffer(StateData, NewState),
+    reroute_buffer(StateData, NewState, Reason),
     SmState#sm_state{buffer = [], buffer_size = 0}.
 
-reroute_buffer(StateData, #sm_state{buffer = Buffer, peer = {gen_statem, {Pid, _}}}) ->
+reroute_buffer(StateData, #sm_state{buffer = Buffer, peer = {gen_statem, {Pid, _}}}, _Reason) ->
+    mongoose_c2s:reroute_buffer_to_pid(StateData, Pid, Buffer);
+reroute_buffer(StateData, #sm_state{buffer = Buffer}, {shutdown, {replaced, Pid}}) ->
     mongoose_c2s:reroute_buffer_to_pid(StateData, Pid, Buffer);
-reroute_buffer(StateData, #sm_state{buffer = Buffer}) ->
+reroute_buffer(StateData, #sm_state{buffer = Buffer}, _Reason) ->
     mongoose_c2s:reroute_buffer(StateData, Buffer).
 
 add_delay_elements_to_buffer(#sm_state{buffer = Buffer} = SmState, FromServer) ->
@@ -410,7 +429,7 @@ add_delay_elements_to_buffer(#sm_state{buffer = Buffer} = SmState, FromServer) -
 terminate(Reason, C2SState, StateData) ->
     ?LOG_DEBUG(#{what => stream_mgmt_statem_terminate, reason => Reason,
                  c2s_state => C2SState, c2s_data => StateData}),
-    mongoose_c2s:terminate({shutdown, ?MODULE}, C2SState, StateData).
+    mongoose_c2s:terminate(Reason, C2SState, StateData).
 
 -spec handle_stream_mgmt(mongoose_acc:t(), mongoose_c2s_hooks:params(), exml:element()) ->
     mongoose_c2s_hooks:result().
diff --git a/test/ejabberd_sm_SUITE.erl b/test/ejabberd_sm_SUITE.erl
@@ -366,7 +366,7 @@ too_many_sessions(_C) ->
     [given_session_opened(Sid, USR) || {Sid, USR} <- UserSessions],
 
     receive
-        {forwarded, _Sid, {'$gen_cast', {exit, <<"Replaced by new connection">>}}} ->
+        {forwarded, _Sid, {'$gen_cast', {exit, {replaced, _Pid}}}} ->
             ok;
         Message ->
             ct:fail("Unexpected message: ~p", [Message])
