Bring a bit of symmetry between both async mam workers

esl · Nov 28, 2021 · 928eb74 · 928eb74
1 parent ec9d240
commit 928eb74
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 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,6 +21,7 @@
          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,
@@ -106,6 +107,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