MAM shaper touches

src/ejabberd_sup.erl

@@ -146,9 +146,9 @@ init([]) ->
          worker,
          [mod_muc_iq]},
     ShaperSup =
-        {ejabberd_shaper_sup,
-          {ejabberd_shaper_sup, start_link, []},
-          permanent, infinity, supervisor, [ejabberd_shaper_sup]},
+        {mongoose_shaper_sup,
+          {mongoose_shaper_sup, start_link, []},
+          permanent, infinity, supervisor, [mongoose_shaper_sup]},
     PG =
         {pg,
           {pg, start_link, [mim_scope]},

src/mam/mod_mam_utils.erl

@@ -105,12 +105,8 @@
 
 -ifdef(MAM_INLINE_UTILS).
 -compile({inline, [
-                   rsm_ns_binary/0,
-                   mam_ns_binary/0,
-                   is_archived_elem_for/2,
                    is_valid_message/4,
                    is_valid_message_type/3,
-                   is_valid_message_children/3,
                    encode_compact_uuid/2,
                    get_one_of_path/3,
                    delay/2,
@@ -135,10 +131,6 @@
 
 -export_type([direction/0, retraction_id/0, retraction_info/0]).
 
-%% Constants
-rsm_ns_binary() -> <<"http://jabber.org/protocol/rsm">>.
-
-
 %% ----------------------------------------------------------------------
 %% Datetime types
 -type ne_binary() :: <<_:8, _:_*8>>.
@@ -510,7 +502,7 @@ result_set(FirstId, LastId, undefined, undefined)
               || LastId =/= undefined],
     #xmlel{
        name = <<"set">>,
-       attrs = [{<<"xmlns">>, rsm_ns_binary()}],
+       attrs = [{<<"xmlns">>, ?NS_RSM}],
        children = FirstEl ++ LastEl};
 result_set(FirstId, LastId, FirstIndexI, CountI)
   when ?MAYBE_BIN(FirstId), ?MAYBE_BIN(LastId) ->
@@ -528,7 +520,7 @@ result_set(FirstId, LastId, FirstIndexI, CountI)
                  children = [#xmlcdata{content = integer_to_binary(CountI)}]},
     #xmlel{
        name = <<"set">>,
-       attrs = [{<<"xmlns">>, rsm_ns_binary()}],
+       attrs = [{<<"xmlns">>, ?NS_RSM}],
        children = FirstEl ++ LastEl ++ [CountEl]}.
 
 
@@ -1111,14 +1103,13 @@ maybe_set_client_xmlns(true, Packet) ->
 maybe_set_client_xmlns(false, Packet) ->
     Packet.
 
-
 -spec action_to_shaper_name(mam_iq:action()) -> atom().
 action_to_shaper_name(Action) ->
     list_to_atom(atom_to_list(Action) ++ "_shaper").
 
 -spec action_to_global_shaper_name(mam_iq:action()) -> atom().
-action_to_global_shaper_name(Action) -> list_to_atom(atom_to_list(Action) ++ "_global_shaper").
-
+action_to_global_shaper_name(Action) ->
+    list_to_atom(atom_to_list(Action) ++ "_global_shaper").
 
 -spec wait_shaper(mongooseim:host_type(), jid:server(), mam_iq:action(), jid:jid()) ->
     'ok' | {'error', 'max_delay_reached'}.

src/mongoose_shaper_sup.erl

@@ -0,0 +1,58 @@
+-module(mongoose_shaper_sup).
+
+-behaviour(supervisor).
+
+%% API
+-export([start_link/0, get_workers/0, select_worker/1]).
+-ignore_xref([start_link/0]).
+
+%% Supervisor callbacks
+-export([init/1]).
+
+-spec start_link() -> {ok, pid()} | ignore | {error, term()}.
+start_link() ->
+    supervisor:start_link({local, ?MODULE}, ?MODULE, []).
+
+-spec init([]) -> {ok, {#{strategy => one_for_one, intensity => 100, period => 5},
+                        [supervisor:child_spec()]}}.
+init([]) ->
+    SupFlags = #{strategy => one_for_one, intensity => 100, period => 5},
+    WorkerNames = worker_names(),
+    TupleNames = list_to_tuple(WorkerNames),
+    persistent_term:put(?MODULE, TupleNames),
+    Shapers = [child_spec(Name) || Name <- WorkerNames],
+    {ok, { SupFlags, Shapers }}.
+
+-spec child_spec(atom()) -> supervisor:child_spec().
+child_spec(ProcName) ->
+    #{id => ProcName,
+      start => {shaper_srv, start_link, [ProcName]},
+      restart => permanent,
+      shutdown => 5000,
+      type => worker,
+      modules => [shaper_srv]}.
+
+-spec get_workers() -> [atom()].
+get_workers() ->
+    tuple_to_list(persistent_term:get(?MODULE)).
+
+-spec select_worker(term()) -> atom().
+select_worker(Key) ->
+    N = 1 + erlang:phash2(Key, worker_count()),
+    Workers = persistent_term:get(?MODULE),
+    element(N, Workers).
+
+-spec worker_names() -> [atom()].
+worker_names() ->
+    [build_worker_name(N) || N <- lists:seq(1, worker_count())].
+
+-spec build_worker_name(integer()) -> atom().
+build_worker_name(N) ->
+    list_to_atom(worker_prefix() ++ integer_to_list(N)).
+
+-spec worker_prefix() -> string().
+worker_prefix() ->
+    "mongoose_shaper_".
+
+worker_count() ->
+    10.