diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl index 42aa0b50c0..ce93093a32 100644 --- a/src/ejabberd_sm.erl +++ b/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() -> diff --git a/test/ejabberd_sm_SUITE.erl b/test/ejabberd_sm_SUITE.erl index fdce27ea21..7cbbb2fd35 100644 --- a/test/ejabberd_sm_SUITE.erl +++ b/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,6 +88,10 @@ 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"}. @@ -93,12 +99,17 @@ init_redis_group(_, _) -> 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.