Merge pull request #4146 from esl/cets/mod_register

Cets/mod register

src/ejabberd_sup.erl

@@ -30,6 +30,7 @@
 
 -export([start_link/0, init/1]).
 -export([start_child/1, start_child/2, stop_child/1]).
+-export([create_ets_table/2]).
 
 -include("mongoose_logger.hrl").
 
@@ -119,3 +120,30 @@ worker_spec(Mod) ->
 
 worker_spec(Mod, Args) ->
     {Mod, {Mod, start_link, Args}, permanent, timer:seconds(5), worker, [Mod]}.
+
+-spec create_ets_table(atom(), list()) -> ok.
+create_ets_table(TableName, TableOpts) ->
+    case does_table_exist(TableName) of
+        true -> ok;
+        false ->
+            Opts = maybe_add_heir(whereis(?MODULE), self(), TableOpts),
+            ets:new(TableName, Opts),
+            ok
+    end.
+
+does_table_exist(TableName) ->
+    undefined =/= ets:info(TableName, name).
+
+%% In tests or when module is started in run-time, we need to set heir to the
+%% ETS table, otherwise it will be destroyed when the creator's process finishes.
+%% When started normally during node start up, self() =:= EjdSupPid and there
+%% is no need for setting heir
+maybe_add_heir(EjdSupPid, EjdSupPid, BaseOpts) when is_pid(EjdSupPid) ->
+    BaseOpts;
+maybe_add_heir(EjdSupPid, _Self, BaseOpts) when is_pid(EjdSupPid) ->
+    case lists:keymember(heir, 1, BaseOpts) of
+        true -> BaseOpts;
+        false -> [{heir, EjdSupPid, testing} | BaseOpts]
+    end;
+maybe_add_heir(_, _, BaseOpts) ->
+    BaseOpts.

src/global_distrib/mod_global_distrib_utils.erl

@@ -111,14 +111,7 @@ create_ets(Names) ->
 create_ets(Names, Type) when is_list(Names) ->
     [create_ets(Name, Type) || Name <- Names];
 create_ets(Name, Type) ->
-    Self = self(),
-    Heir = case whereis(ejabberd_sup) of
-               undefined -> {heir, none};
-               Self -> {heir, none};
-               Pid -> {heir, Pid, testing}
-           end,
-
-    ets:new(Name, [named_table, public, Type, {read_concurrency, true}, Heir]).
+    ejabberd_sup:create_ets_table(Name, [named_table, public, Type, {read_concurrency, true}]).
 
 -spec resolve_endpoints([{inet:ip_address() | string(), inet:port_number()}]) -> [endpoint()].
 resolve_endpoints(Endpoints) when is_list(Endpoints) ->

src/mod_keystore.erl

@@ -49,7 +49,7 @@
 
 -spec start(mongooseim:host_type(), gen_mod:module_opts()) -> ok.
 start(HostType, Opts) ->
-    create_keystore_ets(),
+    ejabberd_sup:create_ets_table(keystore, [named_table, public, {read_concurrency, true}]),
     mod_keystore_backend:init(HostType, Opts),
     init_keys(HostType, Opts),
     ok.
@@ -129,36 +129,10 @@ get_key(HandlerAcc, #{key_id := KeyID}, _) ->
 %%
 %% Internal functions
 %%
-
-create_keystore_ets() ->
-    case does_table_exist(keystore) of
-        true -> ok;
-        false ->
-            BaseOpts = [named_table, public,
-                        {read_concurrency, true}],
-            Opts = maybe_add_heir(whereis(ejabberd_sup), self(), BaseOpts),
-            ets:new(keystore, Opts),
-            ok
-    end.
-
-%% In tests or when module is started in run-time, we need to set heir to the
-%% ETS table, otherwise it will be destroy when the creator's process finishes.
-%% When started normally during node start up, self() =:= EjdSupPid and there
-%% is no need for setting heir
-maybe_add_heir(EjdSupPid, EjdSupPid, BaseOpts) when is_pid(EjdSupPid) ->
-     BaseOpts;
-maybe_add_heir(EjdSupPid, _Self, BaseOpts) when is_pid(EjdSupPid) ->
-      [{heir, EjdSupPid, testing} | BaseOpts];
-maybe_add_heir(_, _, BaseOpts) ->
-      BaseOpts.
-
 clear_keystore_ets(HostType) ->
     Pattern = {{'_', HostType}, '$1'},
     ets:match_delete(keystore, Pattern).
 
-does_table_exist(NameOrTID) ->
-    ets:info(NameOrTID, name) /= undefined.
-
 init_keys(HostType, Opts = #{keys := Keys}) ->
     maps:map(fun(KeyName, KeyType) -> init_key({KeyName, KeyType}, HostType, Opts) end, Keys).
 

src/mod_register.erl

@@ -46,18 +46,17 @@
 -include("mongoose.hrl").
 -include("jlib.hrl").
 -include("mongoose_config_spec.hrl").
+-define(TABLE, mod_register_ip).
 
 -spec start(mongooseim:host_type(), gen_mod:module_opts()) -> ok.
 start(HostType, #{iqdisc := IQDisc}) ->
     [gen_iq_handler:add_iq_handler_for_domain(HostType, ?NS_REGISTER, Component, Fn, #{}, IQDisc) ||
         {Component, Fn} <- iq_handlers()],
-    %% TODO Add CETS backend, use mongoose_mnesia
-    mnesia:create_table(mod_register_ip,
-        [{ram_copies, [node()]},
-         {local_content, true},
-         {attributes, [key, value]}]),
-    mnesia:add_table_copy(mod_register_ip, node(), ram_copies),
-    ok.
+    Concurrency = case IQDisc of
+                      one_queue -> [];
+                      _ -> [{read_concurrency, true}]
+                  end,
+    ejabberd_sup:create_ets_table(?TABLE, [named_table, public | Concurrency]).
 
 -spec stop(mongooseim:host_type()) -> ok.
 stop(HostType) ->
@@ -66,7 +65,8 @@ stop(HostType) ->
     ok.
 
 iq_handlers() ->
-    [{ejabberd_local, fun ?MODULE:process_iq/5}, {ejabberd_sm, fun ?MODULE:process_iq/5}].
+    [{ejabberd_local, fun ?MODULE:process_iq/5},
+     {ejabberd_sm, fun ?MODULE:process_iq/5}].
 
 -spec hooks(mongooseim:host_type()) -> gen_hook:hook_list().
 hooks(HostType) ->
@@ -438,51 +438,24 @@ check_timeout(Source) ->
     Timeout = mongoose_config:get_opt(registration_timeout),
     case is_integer(Timeout) of
         true ->
-            Priority = -(erlang:system_time(second)),
-            CleanPriority = Priority + Timeout,
-            F = fun() -> check_and_store_ip_entry(Source, Priority, CleanPriority) end,
-
-            case mnesia:transaction(F) of
-                {atomic, Res} ->
-                    Res;
-                {aborted, Reason} ->
-                    ?LOG_ERROR(#{what => reg_check_timeout_failed,
-                                 reg_source => Source, reason => Reason}),
-                    true
-            end;
+            TS = erlang:system_time(second),
+            clean_ets(TS - Timeout),
+            check_and_store_ip_entry(Source, TS);
         false ->
             true
     end.
 
-check_and_store_ip_entry(Source, Priority, CleanPriority) ->
-    Treap = case mnesia:read(mod_register_ip, treap, write) of
-                [] ->
-                    treap:empty();
-                [{mod_register_ip, treap, T}] -> T
-            end,
-    Treap1 = clean_treap(Treap, CleanPriority),
-    case treap:lookup(Source, Treap1) of
-        error ->
-            Treap2 = treap:insert(Source, Priority, [],
-                                  Treap1),
-            mnesia:write({mod_register_ip, treap, Treap2}),
+check_and_store_ip_entry(Source, Timestamp) ->
+    case ets:member(?TABLE, Source) of
+        false ->
+            ets:insert(?TABLE, {Source, Timestamp}),
             true;
-        {ok, _, _} ->
-            mnesia:write({mod_register_ip, treap, Treap1}),
+        true ->
             false
     end.
 
-clean_treap(Treap, CleanPriority) ->
-    case treap:is_empty(Treap) of
-        true ->
-            Treap;
-        false ->
-            {_Key, Priority, _Value} = treap:get_root(Treap),
-            case Priority > CleanPriority of
-                true -> clean_treap(treap:delete_root(Treap), CleanPriority);
-                false -> Treap
-            end
-    end.
+clean_ets(CleanTimestamp) ->
+    ets:select_delete(?TABLE, [{ {'_', '$1'}, [{'<', '$1', CleanTimestamp}], [true]}]).
 
 ip_to_string(Source) when is_tuple(Source) -> inet_parse:ntoa(Source);
 ip_to_string(undefined) -> "undefined";

src/wpool/mongoose_wpool.erl

@@ -106,17 +106,12 @@ ensure_started() ->
         _ ->
             ok
     end,
-
-    case ets:info(?MODULE) of
-        undefined ->
-            % we set heir here because the whole thing may be started by an ephemeral process
-            ets:new(?MODULE, [named_table, public,
-                {read_concurrency, true},
-                {keypos, #mongoose_wpool.name},
-                {heir, whereis(mongoose_wpool_sup), undefined}]);
-        _ ->
-            ok
-    end.
+    ejabberd_sup:create_ets_table(
+      ?MODULE,
+      [named_table, public,
+       {read_concurrency, true},
+       {keypos, #mongoose_wpool.name},
+       {heir, whereis(mongoose_wpool_sup), undefined}]).
 
 start_configured_pools() ->
     Pools = mongoose_config:get_opt(outgoing_pools),

src/wpool/mongoose_wpool_http.erl

@@ -20,18 +20,7 @@
 %% mongoose_wpool callbacks
 -spec init() -> ok.
 init() ->
-    case ets:info(?MODULE) of
-        undefined ->
-            Heir = case whereis(ejabberd_sup) of
-                       undefined -> [];
-                       Pid -> [{heir, Pid, undefined}]
-                   end,
-            ets:new(?MODULE,
-                    [named_table, public, {read_concurrency, true} | Heir]),
-            ok;
-        _ ->
-            ok
-    end.
+    ejabberd_sup:create_ets_table(?MODULE, [named_table, public, {read_concurrency, true}]).
 
 -spec start(mongooseim:host_type_or_global(), mongoose_wpool:tag(),
             mongoose_wpool:pool_opts(), mongoose_wpool:conn_opts()) -> {ok, pid()} | {error, any()}.

src/wpool/mongoose_wpool_rdbms.erl

@@ -9,18 +9,8 @@
 %% mongoose_wpool callbacks
 -spec init() -> ok.
 init() ->
-    case ets:info(prepared_statements) of
-        undefined ->
-            Heir = case whereis(ejabberd_sup) of
-                       undefined -> [];
-                       Pid -> [{heir, Pid, undefined}]
-                   end,
-            ets:new(prepared_statements,
-                    [named_table, public, {read_concurrency, true} | Heir]),
-            ok;
-        _ ->
-            ok
-    end.
+    ejabberd_sup:create_ets_table(
+      prepared_statements, [named_table, public, {read_concurrency, true}]).
 
 -spec start(mongooseim:host_type_or_global(), mongoose_wpool:tag(),
             mongoose_wpool:pool_opts(), mongoose_wpool:conn_opts()) -> {ok, pid()} | {error, any()}.