Migrate unique count cache to internal DB

src/ejabberd_sm.erl

@@ -87,7 +87,8 @@
 -export([do_route/4]).
 
 -ignore_xref([do_filter/3, do_route/4, get_unique_sessions_number/0,
-              get_user_present_pids/2, start_link/0, user_resources/2, sm_backend/0]).
+              get_user_present_pids/2, start_link/0, user_resources/2, sm_backend/0,
+              stop_probes/0, start_probes/0]).
 
 -include("mongoose.hrl").
 -include("jlib.hrl").
@@ -111,6 +112,7 @@
 -type backend() :: ejabberd_sm_mnesia | ejabberd_sm_redis | ejabberd_sm_cets.
 -type close_reason() :: resumed | normal | replaced.
 -type info_key() :: atom().
+-type internal_db() :: mnesia | cets.
 
 -export_type([session/0,
               sid/0,
@@ -313,7 +315,15 @@ get_session_pid(JID) ->
 
 -spec get_unique_sessions_number() -> integer().
 get_unique_sessions_number() ->
-    ejabberd_sm_backend:unique_count().
+    InternalDB = get_internal_db(),
+    try
+        C = ejabberd_sm_backend:unique_count(),
+        set_cached_unique_count(C, InternalDB),
+        C
+    catch
+        _:_ ->
+            get_cached_unique_count(InternalDB)
+    end.
 
 -spec get_total_sessions_number() -> integer().
 get_total_sessions_number() ->
@@ -478,11 +488,22 @@ init([]) ->
                   ?ALL_HOST_TYPES),
     %% Create metrics after backend has started, otherwise probe could have null value
     start_probes(),
+    init_cache(),
     {ok, #state{}}.
 
 start_probes() ->
     mongoose_instrument:set_up(instrumentation(global)).
 
+init_cache() ->
+    init_cache(get_internal_db()).
+
+init_cache(mnesia) ->
+    mongoose_mnesia:create_table(unique_sessions_table, [{attributes, [key, count]},
+                                                         {ram_copies, [node()]}]);
+init_cache(cets) ->
+    cets:start(cets_unique_sessions, #{}),
+    cets_discovery:add_table(mongoose_cets_discovery, cets_unique_sessions).
+
 -spec hooks(binary()) -> [gen_hook:hook_tuple()].
 hooks(HostType) ->
     [
@@ -973,6 +994,40 @@ user_resources(UserStr, ServerStr) ->
     Resources = get_user_resources(JID),
     lists:sort(Resources).
 
+-spec get_cached_unique_count(internal_db()) -> non_neg_integer().
+get_cached_unique_count(mnesia) ->
+    F = fun() ->
+        case mnesia:read({unique_sessions_table, unique_sessions_count}) of
+            [{unique_sessions_table, unique_sessions_count, C}] -> C;
+            [] -> 0
+        end
+    end,
+    {atomic, Result} = mnesia:transaction(F),
+    Result;
+get_cached_unique_count(cets) ->
+    case ets:lookup(cets_unique_sessions, unique_sessions_count) of
+        [{unique_sessions_count, C}] -> C;
+        [] -> 0
+    end.
+
+-spec set_cached_unique_count(non_neg_integer(), internal_db()) -> ok | {error, any()}.
+set_cached_unique_count(C, mnesia) ->
+    mnesia:transaction(
+        fun() ->
+            mnesia:write({unique_sessions_table, unique_sessions_count, C})
+        end);
+set_cached_unique_count(C, cets) ->
+    cets:insert_serial(cets_unique_sessions, {unique_sessions_count, C}).
+
+-spec get_internal_db() -> internal_db().
+get_internal_db() ->
+    case mongoose_config:get_opt(internal_databases) of
+        #{mnesia := _} ->
+            mnesia;
+        #{cets := _} ->
+            cets
+    end.
+
 %% It is used from big tests
 -spec sm_backend() -> backend().
 sm_backend() ->

test/ejabberd_sm_SUITE.erl

@@ -66,12 +66,14 @@ tests() ->
 init_per_group(mnesia, Config) ->
     ok = mnesia:create_schema([node()]),
     ok = mnesia:start(),
+    mongoose_config:set_opt(internal_databases, #{mnesia => #{}}),
     [{backend, ejabberd_sm_mnesia} | Config];
 init_per_group(redis, Config) ->
     init_redis_group(is_redis_running(), Config);
 init_per_group(cets, Config) ->
     DiscoOpts = #{name => mongoose_cets_discovery, disco_file => "does_not_exist.txt"},
     {ok, Pid} = cets_discovery:start(DiscoOpts),
+    mongoose_config:set_opt(internal_databases, #{cets => #{}}),
     [{backend, ejabberd_sm_cets}, {cets_disco_pid, Pid} | Config].
 
 init_redis_group(true, Config) ->
@@ -86,19 +88,28 @@ init_redis_group(true, Config) ->
                   receive stop -> ok end
           end),
     receive ready -> ok after timer:seconds(30) -> ct:fail(test_helper_not_ready) end,
+    %% Set up mnesia for unique sessions cache
+    ok = mnesia:create_schema([node()]),
+    ok = mnesia:start(),
+    mongoose_config:set_opt(internal_databases, #{mnesia => #{}}),
     [{backend, ejabberd_sm_redis} | Config];
 init_redis_group(_, _) ->
     {skip, "redis not running"}.
 
 end_per_group(mnesia, Config) ->
     mnesia:stop(),
     mnesia:delete_schema([node()]),
+    mongoose_config:unset_opt(internal_databases),
     Config;
 end_per_group(cets, Config) ->
     exit(proplists:get_value(cets_disco_pid, Config), kill),
+    mongoose_config:unset_opt(internal_databases),
     Config;
 end_per_group(redis, Config) ->
     whereis(test_helper) ! stop,
+    mnesia:stop(),
+    mnesia:delete_schema([node()]),
+    mongoose_config:unset_opt(internal_databases),
     Config.
 
 init_per_testcase(too_many_sessions, Config) ->
@@ -394,8 +405,8 @@ unique_count_while_removing_entries(C) ->
     [given_session_opened(Sid, USR) || {Sid, USR} <- UsersWithManyResources],
     set_test_case_meck_unique_count_crash(?B(C)),
     USDict = get_unique_us_dict(UsersWithManyResources),
-    %% Check if unique count equals prev cached value
     UniqueCount = ejabberd_sm:get_unique_sessions_number(),
+    % ct:fail("get_unique_sessions_number/0 should have failed"),
     meck:unload(?B(C)),
     true = UniqueCount /= dict:size(USDict) + UniqueCount.