Rewrite mnesia counter functions to use dirty_update_counter (EJAB-1177)

(thanks to Juan Pablo Carlino and Alexey Shchepin)

src/ejabberd_auth_internal.erl

@@ -72,7 +72,11 @@ update_reg_users_counter_table(Server) ->
     Set = get_vh_registered_users(Server),
     Size = length(Set),
     LServer = exmpp_jid:prep_domain(exmpp_jid:parse(Server)),
-    set_vh_registered_users_counter(LServer, Size).
+    F = fun() ->
+        mnesia:write(#reg_users_counter{vhost = LServer,
+					count = Size})
+    end,
+    mnesia:sync_dirty(F).
 
 %% @spec () -> bool()
 
@@ -162,7 +166,9 @@ try_register(User, Server, Password) ->
 			    [] ->
 				mnesia:write(#passwd{us = US,
 						     password = Password}),
-				inc_vh_registered_users_counter(LServer),
+				mnesia:dirty_update_counter(
+						    reg_users_counter,
+						    exmpp_jid:prep_domain(exmpp_jid:parse(Server)), 1),
 				ok;
 			    [_E] ->
 				exists
@@ -285,40 +291,6 @@ get_vh_registered_users_number(Server, [{prefix, Prefix}]) when is_list(Prefix)
 get_vh_registered_users_number(Server, _) ->
     get_vh_registered_users_number(Server).
 
-inc_vh_registered_users_counter(LServer) ->
-    F = fun() ->
-		case mnesia:wread({reg_users_counter, LServer}) of
-		    [C] ->
-			Count = C#reg_users_counter.count + 1,
-			C2 = C#reg_users_counter{count = Count},
-			mnesia:write(C2);
-		    _ ->
-			mnesia:write(#reg_users_counter{vhost = LServer,
-						      count = 1})
-		end
-	end,
-    mnesia:sync_dirty(F).
-
-dec_vh_registered_users_counter(LServer) ->
-    F = fun() ->
-		case mnesia:wread({reg_users_counter, LServer}) of
-		    [C] ->
-			Count = C#reg_users_counter.count - 1,
-			C2 = C#reg_users_counter{count = Count},
-			mnesia:write(C2);
-		    _ ->
-			error
-		end
-	end,
-    mnesia:sync_dirty(F).
-
-set_vh_registered_users_counter(LServer, Count) ->
-    F = fun() ->
-		mnesia:write(#reg_users_counter{vhost = LServer,
-						count = Count})
-	end,
-    mnesia:sync_dirty(F).
-
 %% @spec (User, Server) -> Password | false
 %%     User = string()
 %%     Server = string()
@@ -396,7 +368,8 @@ remove_user(User, Server) ->
 	US = {LUser, LServer},
 	F = fun() ->
 		    mnesia:delete({passwd, US}),
-		    dec_vh_registered_users_counter(LServer)
+		    mnesia:dirty_update_counter(reg_users_counter,
+						exmpp_jid:prep_domain(exmpp_jid:parse(Server)), -1)
 	    end,
 	mnesia:transaction(F),
 	ok
@@ -420,7 +393,8 @@ remove_user(User, Server, Password) ->
 		    case mnesia:read({passwd, US}) of
 			[#passwd{password = Password}] ->
 			    mnesia:delete({passwd, US}),
-			    dec_vh_registered_users_counter(LServer),
+			    mnesia:dirty_update_counter(reg_users_counter,
+							exmpp_jid:prep_domain(exmpp_jid:parse(Server)), -1),
 			    ok;
 			[_] ->
 			    not_allowed;

src/ejabberd_sm.erl

@@ -113,7 +113,8 @@ route(From, To, Packet) ->
 
 open_session(SID, JID, Info) when ?IS_JID(JID) ->
     set_session(SID, JID, undefined, Info),
-    inc_session_counter(exmpp_jid:domain(JID)),
+    mnesia:dirty_update_counter(session_counter,
+				exmpp_jid:domain(JID), 1),
     check_for_sessions_to_replace(JID),
     ejabberd_hooks:run(sm_register_connection_hook, exmpp_jid:prep_domain(JID),
 		       [SID, JID, Info]).
@@ -125,7 +126,8 @@ close_session(SID, JID ) when ?IS_JID(JID)->
     end,
     F = fun() ->
 		mnesia:delete({session, SID}),
-		dec_session_counter(exmpp_jid:domain(JID))
+		mnesia:dirty_update_counter(session_counter,
+					    exmpp_jid:domain(JID), -1)
 	end,
     mnesia:sync_dirty(F),
     ejabberd_hooks:run(sm_remove_connection_hook, exmpp_jid:prep_domain(JID),
@@ -299,6 +301,7 @@ init([]) ->
     mnesia:add_table_index(session, usr),
     mnesia:add_table_index(session, us),
     mnesia:add_table_copy(session, node(), ram_copies),
+    mnesia:add_table_copy(session_counter, node(), ram_copies),
     mnesia:subscribe(system),
     ets:new(sm_iqtable, [named_table]),
     lists:foreach(
@@ -364,7 +367,7 @@ handle_info({route, From, To, Packet}, State) ->
     end,
     {noreply, State};
 handle_info({mnesia_system_event, {mnesia_down, Node}}, State) ->
-    clean_table_from_bad_node(Node),
+    recount_session_table(Node),
     {noreply, State};
 handle_info({register_iq_handler, Host, XMLNS, Module, Function}, State) ->
     ets:insert(sm_iqtable, {{XMLNS, Host}, Module, Function}),
@@ -420,7 +423,9 @@ set_session(SID, JID, Priority, Info) ->
 	end,
     mnesia:sync_dirty(F).
 
-clean_table_from_bad_node(Node) ->
+%% Recalculates alive sessions when Node goes down 
+%% and updates session and session_counter tables
+recount_session_table(Node) ->
     F = fun() ->
 		Es = mnesia:select(
 		       session,
@@ -428,12 +433,22 @@ clean_table_from_bad_node(Node) ->
 			 [{'==', {node, '$1'}, Node}],
 			 ['$_']}]),
 		lists:foreach(fun(E) ->
-				      {_, LServer} = E#session.us,
-				      dec_session_counter(LServer),
 				      mnesia:delete({session, E#session.sid})
-			      end, Es)
+			      end, Es),
+		%% reset session_counter table with active sessions
+		mnesia:clear_table(session_counter),
+		lists:foreach(fun(Server) ->
+				LServer = exmpp_jid:prep_domain(exmpp_jid:parse(Server)),
+				Hs = mnesia:select(session,
+				    [{#session{usr = '$1', _ = '_'},
+				    [{'==', {element, 2, '$1'}, LServer}],
+				    ['$1']}]),
+				mnesia:write(
+				    #session_counter{vhost = LServer, 
+						     count = length(Hs)})
+			      end, ?MYHOSTS)
 	end,
-    mnesia:sync_dirty(F).
+    mnesia:async_dirty(F).
 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 
@@ -708,34 +723,6 @@ get_max_user_sessions(JID) ->
 	_ -> ?MAX_USER_SESSIONS
     end.
 
-inc_session_counter(LServer) ->
-    F = fun() ->
-		case mnesia:wread({session_counter, LServer}) of
-		    [C] ->
-			Count = C#session_counter.count + 1,
-			C2 = C#session_counter{count = Count},
-			mnesia:write(C2);
-		    _ ->
-			mnesia:write(#session_counter{vhost = LServer,
-						      count = 1})
-		end
-	end,
-    mnesia:sync_dirty(F).
-
-dec_session_counter(LServer) ->
-    F = fun() ->
-		case mnesia:wread({session_counter, LServer}) of
-		    [C] ->
-			Count = C#session_counter.count - 1,
-			C2 = C#session_counter{count = Count},
-			mnesia:write(C2);
-		    _ ->
-			error
-		end
-	end,
-    mnesia:sync_dirty(F).
-
-
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 
 process_iq(From, To, Packet) ->