Implement cache for mod_last

src/mod_last.erl

@@ -45,17 +45,24 @@
 -include("mod_privacy.hrl").
 -include("mod_last.hrl").
 
+-define(LAST_CACHE, last_activity_cache).
+
 -callback init(binary(), gen_mod:opts()) -> any().
 -callback import(binary(), #last_activity{}) -> ok | pass.
 -callback get_last(binary(), binary()) ->
-    {ok, non_neg_integer(), binary()} | not_found | {error, any()}.
--callback store_last_info(binary(), binary(), non_neg_integer(), binary()) -> any().
+    {ok, {non_neg_integer(), binary()}} | error | {error, any()}.
+-callback store_last_info(binary(), binary(), non_neg_integer(), binary()) -> ok | {error, any()}.
 -callback remove_user(binary(), binary()) -> any().
+-callback use_cache(binary()) -> boolean().
+-callback cache_nodes(binary()) -> [node()].
+
+-optional_callbacks([use_cache/1, cache_nodes/1]).
 
 start(Host, Opts) ->
     IQDisc = gen_mod:get_opt(iqdisc, Opts, gen_iq_handler:iqdisc(Host)),
     Mod = gen_mod:db_mod(Host, Opts, ?MODULE),
     Mod:init(Host, Opts),
+    init_cache(Mod, Host, Opts),
     gen_iq_handler:add_iq_handler(ejabberd_local, Host,
 				  ?NS_LAST, ?MODULE, process_local_iq, IQDisc),
     gen_iq_handler:add_iq_handler(ejabberd_sm, Host,
@@ -91,6 +98,7 @@ reload(Host, NewOpts, OldOpts) ->
        true ->
 	    ok
     end,
+    init_cache(NewMod, Host, NewOpts),
     case gen_mod:is_equal_opt(iqdisc, NewOpts, OldOpts, gen_iq_handler:iqdisc(Host)) of
 	{false, IQDisc, _} ->
 	    gen_iq_handler:add_iq_handler(ejabberd_local, Host, ?NS_LAST,
@@ -180,13 +188,23 @@ privacy_check_packet(allow, C2SState,
 privacy_check_packet(Acc, _, _, _) ->
     Acc.
 
-%% @spec (LUser::string(), LServer::string()) ->
-%%      {ok, TimeStamp::integer(), Status::string()} | not_found | {error, Reason}
 -spec get_last(binary(), binary()) -> {ok, non_neg_integer(), binary()} |
 				      not_found | {error, any()}.
 get_last(LUser, LServer) ->
     Mod = gen_mod:db_mod(LServer, ?MODULE),
-    Mod:get_last(LUser, LServer).
+    Res = case use_cache(Mod, LServer) of
+	      true ->
+		  ets_cache:lookup(
+		    ?LAST_CACHE, {LUser, LServer},
+		    fun() -> Mod:get_last(LUser, LServer) end);
+	      false ->
+		  Mod:get_last(LUser, LServer)
+	  end,
+    case Res of
+	{ok, {TimeStamp, Status}} -> {ok, TimeStamp, Status};
+	error -> not_found;
+	Err -> Err
+    end.
 
 -spec get_last_iq(iq(), binary(), binary()) -> iq().
 get_last_iq(#iq{lang = Lang} = IQ, LUser, LServer) ->
@@ -226,7 +244,16 @@ store_last_info(User, Server, TimeStamp, Status) ->
     LUser = jid:nodeprep(User),
     LServer = jid:nameprep(Server),
     Mod = gen_mod:db_mod(LServer, ?MODULE),
-    Mod:store_last_info(LUser, LServer, TimeStamp, Status).
+    case use_cache(Mod, LServer) of
+	true ->
+	    ets_cache:update(
+	      ?LAST_CACHE, {LUser, LServer}, {ok, {TimeStamp, Status}},
+	      fun() ->
+		      Mod:store_last_info(LUser, LServer, TimeStamp, Status)
+	      end, cache_nodes(Mod, LServer));
+	false ->
+	    Mod:store_last_info(LUser, LServer, TimeStamp, Status)
+    end.
 
 -spec get_last_info(binary(), binary()) -> {ok, non_neg_integer(), binary()} |
 					   not_found.
@@ -241,7 +268,51 @@ remove_user(User, Server) ->
     LUser = jid:nodeprep(User),
     LServer = jid:nameprep(Server),
     Mod = gen_mod:db_mod(LServer, ?MODULE),
-    Mod:remove_user(LUser, LServer).
+    Mod:remove_user(LUser, LServer),
+    ets_cache:delete(?LAST_CACHE, {LUser, LServer}, cache_nodes(Mod, LServer)).
+
+-spec init_cache(module(), binary(), gen_mod:opts()) -> ok.
+init_cache(Mod, Host, Opts) ->
+    case use_cache(Mod, Host) of
+	true ->
+	    CacheOpts = cache_opts(Host, Opts),
+	    ets_cache:new(?LAST_CACHE, CacheOpts);
+	false ->
+	    ets_cache:delete(?LAST_CACHE)
+    end.
+
+-spec cache_opts(binary(), gen_mod:opts()) -> [proplists:property()].
+cache_opts(Host, Opts) ->
+    MaxSize = gen_mod:get_opt(
+		cache_size, Opts,
+		ejabberd_config:cache_size(Host)),
+    CacheMissed = gen_mod:get_opt(
+		    cache_missed, Opts,
+		    ejabberd_config:cache_missed(Host)),
+    LifeTime = case gen_mod:get_opt(
+		      cache_life_time, Opts,
+		      ejabberd_config:cache_life_time(Host)) of
+		   infinity -> infinity;
+		   I -> timer:seconds(I)
+	       end,
+    [{max_size, MaxSize}, {cache_missed, CacheMissed}, {life_time, LifeTime}].
+
+-spec use_cache(module(), binary()) -> boolean().
+use_cache(Mod, Host) ->
+    case erlang:function_exported(Mod, use_cache, 1) of
+	true -> Mod:use_cache(Host);
+	false ->
+	    gen_mod:get_module_opt(
+	      Host, ?MODULE, use_cache,
+	      ejabberd_config:use_cache(Host))
+    end.
+
+-spec cache_nodes(module(), binary()) -> [node()].
+cache_nodes(Mod, Host) ->
+    case erlang:function_exported(Mod, cache_nodes, 1) of
+	true -> Mod:cache_nodes(Host);
+	false -> ejabberd_cluster:get_nodes()
+    end.
 
 import_info() ->
     [{<<"last">>, 3}].
@@ -270,4 +341,11 @@ depends(_Host, _Opts) ->
 
 mod_opt_type(db_type) -> fun(T) -> ejabberd_config:v_db(?MODULE, T) end;
 mod_opt_type(iqdisc) -> fun gen_iq_handler:check_type/1;
-mod_opt_type(_) -> [db_type, iqdisc].
+mod_opt_type(O) when O == cache_life_time; O == cache_size ->
+    fun (I) when is_integer(I), I > 0 -> I;
+        (infinity) -> infinity
+    end;
+mod_opt_type(O) when O == use_cache; O == cache_missed ->
+    fun (B) when is_boolean(B) -> B end;
+mod_opt_type(_) ->
+    [db_type, iqdisc, cache_life_time, cache_size, use_cache, cache_missed].

src/mod_last_mnesia.erl

@@ -27,7 +27,8 @@
 -behaviour(mod_last).
 
 %% API
--export([init/2, import/2, get_last/2, store_last_info/4, remove_user/2]).
+-export([init/2, import/2, get_last/2, store_last_info/4,
+	 remove_user/2, use_cache/1]).
 -export([need_transform/1, transform/1]).
 
 -include("mod_last.hrl").
@@ -38,31 +39,36 @@
 %%%===================================================================
 init(_Host, _Opts) ->
     ejabberd_mnesia:create(?MODULE, last_activity,
-			   [{disc_copies, [node()]},
+			   [{disc_only_copies, [node()]},
 			    {attributes, record_info(fields, last_activity)}]).
 
+use_cache(Host) ->
+    case mnesia:table_info(last_activity, storage_type) of
+	disc_only_copies ->
+	    gen_mod:get_module_opt(
+	      Host, mod_last, use_cache,
+	      ejabberd_config:use_cache(Host));
+	_ ->
+	    false
+    end.
+
 get_last(LUser, LServer) ->
     case mnesia:dirty_read(last_activity, {LUser, LServer}) of
 	[] ->
-	    not_found;
+	    error;
 	[#last_activity{timestamp = TimeStamp,
 			status = Status}] ->
-	    {ok, TimeStamp, Status}
+	    {ok, {TimeStamp, Status}}
     end.
 
 store_last_info(LUser, LServer, TimeStamp, Status) ->
-    US = {LUser, LServer},
-    F = fun () ->
-		mnesia:write(#last_activity{us = US,
-					    timestamp = TimeStamp,
-					    status = Status})
-	end,
-    mnesia:transaction(F).
+    mnesia:dirty_write(#last_activity{us = {LUser, LServer},
+				      timestamp = TimeStamp,
+				      status = Status}).
 
 remove_user(LUser, LServer) ->
     US = {LUser, LServer},
-    F = fun () -> mnesia:delete({last_activity, US}) end,
-    mnesia:transaction(F).
+    mnesia:dirty_delete({last_activity, US}).
 
 import(_LServer, #last_activity{} = LA) ->
     mnesia:dirty_write(LA).

src/mod_last_riak.erl

@@ -43,19 +43,20 @@ get_last(LUser, LServer) ->
 			   {LUser, LServer}) of
         {ok, #last_activity{timestamp = TimeStamp,
                             status = Status}} ->
-            {ok, TimeStamp, Status};
-        {error, notfound} ->
-            not_found;
-        Err ->
-            Err
+            {ok, {TimeStamp, Status}};
+	{error, notfound} ->
+	    error;
+        _Err ->
+	    %% TODO: log error
+	    {error, db_failure}
     end.
 
 store_last_info(LUser, LServer, TimeStamp, Status) ->
     US = {LUser, LServer},
-    {atomic, ejabberd_riak:put(#last_activity{us = US,
-                                              timestamp = TimeStamp,
-                                              status = Status},
-			       last_activity_schema())}.
+    ejabberd_riak:put(#last_activity{us = US,
+				     timestamp = TimeStamp,
+				     status = Status},
+		      last_activity_schema()).
 
 remove_user(LUser, LServer) ->
     {atomic, ejabberd_riak:delete(last_activity, {LUser, LServer})}.

src/mod_last_sql.erl

@@ -45,17 +45,24 @@ init(_Host, _Opts) ->
 get_last(LUser, LServer) ->
     case catch sql_queries:get_last(LServer, LUser) of
         {selected, []} ->
-            not_found;
+	    error;
         {selected, [{TimeStamp, Status}]} ->
-            {ok, TimeStamp, Status};
+            {ok, {TimeStamp, Status}};
         Reason ->
 	    ?ERROR_MSG("failed to get last for user ~s@~s: ~p",
 		       [LUser, LServer, Reason]),
-	    {error, {invalid_result, Reason}}
+	    {error, db_failure}
     end.
 
 store_last_info(LUser, LServer, TimeStamp, Status) ->
-    sql_queries:set_last_t(LServer, LUser, TimeStamp, Status).
+    case sql_queries:set_last_t(LServer, LUser, TimeStamp, Status) of
+	ok ->
+	    ok;
+	Err ->
+	    ?ERROR_MSG("failed to store last activity for ~s@~s: ~p",
+		       [LUser, LServer, Err]),
+	    {error, db_failure}
+    end.
 
 remove_user(LUser, LServer) ->
     sql_queries:del_last(LServer, LUser).