Bring a bit of symmetry between both async mam workers

esl · Nov 29, 2021 · eee8837 · eee8837
1 parent ec9d240
commit eee8837
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/src/mam/mod_mam_rdbms_async_pool_writer.erl b/src/mam/mod_mam_rdbms_async_pool_writer.erl
@@ -239,7 +239,6 @@ do_run_flush(MessageCount, State = #state{host_type = HostType, max_batch_size =
                     Error -> Error
                 end
         end,
-    [mod_mam_rdbms_arch:retract_message(HostType, Params) || Params <- Acc],
     case InsertResult of
         {updated, _Count} -> ok;
         {error, Reason} ->
@@ -249,7 +248,8 @@ do_run_flush(MessageCount, State = #state{host_type = HostType, max_batch_size =
                          message_count => MessageCount, reason => Reason}),
             ok
     end,
-    mongoose_metrics:update(HostType, modMamFlushed, MessageCount),
+    [mod_mam_rdbms_arch:retract_message(HostType, Params) || Params <- Acc],
+    mongoose_hooks:mam_flush_messages(HostType, MessageCount),
     erlang:garbage_collect(),
     State#state{acc=[], flush_interval_tref=undefined}.
 

diff --git a/src/metrics/mongoose_metrics_mam_hooks.erl b/src/metrics/mongoose_metrics_mam_hooks.erl
@@ -21,14 +21,15 @@
          mam_remove_archive/4,
          mam_lookup_messages/3,
          mam_archive_message/3,
+         mam_flush_messages/3,
          mam_muc_get_prefs/4,
          mam_muc_set_prefs/7,
          mam_muc_remove_archive/4,
          mam_muc_lookup_messages/3,
          mam_muc_archive_message/3,
          mam_muc_flush_messages/3]).
 
--ignore_xref([mam_archive_message/3, mam_get_prefs/4, mam_lookup_messages/3,
+-ignore_xref([mam_archive_message/3, mam_get_prefs/4, mam_lookup_messages/3, mam_flush_messages/3,
               mam_muc_archive_message/3, mam_muc_flush_messages/3, mam_muc_get_prefs/4,
               mam_muc_lookup_messages/3, mam_muc_remove_archive/4, mam_muc_set_prefs/7,
               mam_remove_archive/4, mam_set_prefs/7]).
@@ -47,7 +48,8 @@ get_mam_hooks(Host) ->
         {mam_get_prefs, Host, ?MODULE, mam_get_prefs, 50},
         {mam_archive_message, Host, ?MODULE, mam_archive_message, 50},
         {mam_remove_archive, Host, ?MODULE, mam_remove_archive, 50},
-        {mam_lookup_messages, Host, ?MODULE, mam_lookup_messages, 100}
+        {mam_lookup_messages, Host, ?MODULE, mam_lookup_messages, 100},
+        {mam_flush_messages, Host, ?MODULE, mam_flush_messages, 50}
     ].
 
 %% @doc Here will be declared which hooks should be registered when mod_mam_muc is enabled.
@@ -106,6 +108,10 @@ mam_archive_message(Result, Host, _Params) ->
     mongoose_metrics:update(Host, modMamArchived, 1),
     Result.
 
+mam_flush_messages(Acc, Host, MessageCount) ->
+    mongoose_metrics:update(Host, modMamFlushed, MessageCount),
+    Acc.
+
 %% ----------------------------------------------------------------------------
 %% mod_mam_muc
 

diff --git a/src/mongoose_hooks.erl b/src/mongoose_hooks.erl
@@ -94,7 +94,8 @@
          mam_get_prefs/4,
          mam_remove_archive/3,
          mam_lookup_messages/2,
-         mam_archive_message/2]).
+         mam_archive_message/2,
+         mam_flush_messages/2]).
 
 -export([mam_muc_archive_id/2,
          mam_muc_archive_size/3,
@@ -1008,6 +1009,14 @@ mam_lookup_messages(HookServer, Params) ->
 mam_archive_message(HookServer, Params) ->
     run_hook_for_host_type(mam_archive_message, HookServer, ok, [HookServer, Params]).
 
+%%% @doc The `mam_flush_messages' hook is run after the async bulk write
+%%% happens for messages despite the result of the write.
+-spec mam_flush_messages(HookServer :: jid:lserver(),
+                         MessageCount :: integer()) -> ok.
+mam_flush_messages(HookServer, MessageCount) ->
+    run_hook_for_host_type(mam_flush_messages, HookServer, ok,
+                           [HookServer, MessageCount]).
+
 
 %% MAM MUC related hooks