diff --git a/src/ejabberd_riak_sup.erl b/src/ejabberd_riak_sup.erl index ac02a846f..f5c8f7e4f 100644 --- a/src/ejabberd_riak_sup.erl +++ b/src/ejabberd_riak_sup.erl @@ -77,15 +77,22 @@ is_riak_configured() -> lists:any(fun is_riak_configured/1, ?MYHOSTS). is_riak_configured(Host) -> - ServerConfigured = ejabberd_config:get_option( - {riak_server, Host}, - fun(_) -> true end, false), - PortConfigured = ejabberd_config:get_option( - {riak_port, Host}, - fun(_) -> true end, false), + ServerConfigured = ejabberd_config:has_option({riak_server, Host}), + PortConfigured = ejabberd_config:has_option({riak_port, Host}), + StartIntervalConfigured = ejabberd_config:has_option({riak_start_interval, Host}), + PoolConfigured = ejabberd_config:has_option({riak_pool_size, Host}), + CacertConfigured = ejabberd_config:has_option({riak_cacertfile, Host}), + UserConfigured = ejabberd_config:has_option({riak_username, Host}), + PassConfigured = ejabberd_config:has_option({riak_password, Host}), AuthConfigured = lists:member( ejabberd_auth_riak, ejabberd_auth:auth_modules(Host)), + SMConfigured = ejabberd_config:get_option( + {sm_db_type, Host}, + ejabberd_sm:opt_type(sm_db_type)) == riak, + RouterConfigured = ejabberd_config:get_option( + {router_db_type, Host}, + ejabberd_router:opt_type(router_db_type)) == riak, Modules = ejabberd_config:get_option( {modules, Host}, fun(L) when is_list(L) -> L end, []), @@ -93,7 +100,10 @@ is_riak_configured(Host) -> fun({Module, Opts}) -> gen_mod:db_type(Host, Opts, Module) == riak end, Modules), - ServerConfigured or PortConfigured + ServerConfigured or PortConfigured or StartIntervalConfigured + or PoolConfigured or CacertConfigured + or UserConfigured or PassConfigured + or SMConfigured or RouterConfigured or AuthConfigured or ModuleWithRiakDBConfigured. start_link() -> diff --git a/src/ejabberd_sm_riak.erl b/src/ejabberd_sm_riak.erl new file mode 100644 index 000000000..be6b49029 --- /dev/null +++ b/src/ejabberd_sm_riak.erl @@ -0,0 +1,72 @@ +%%%------------------------------------------------------------------- +%%% @author Evgeny Khramtsov +%%% Created : 15 Apr 2017 by Evgeny Khramtsov +%%% +%%% +%%% ejabberd, Copyright (C) 2002-2017 ProcessOne +%%% +%%% This program is free software; you can redistribute it and/or +%%% modify it under the terms of the GNU General Public License as +%%% published by the Free Software Foundation; either version 2 of the +%%% License, or (at your option) any later version. +%%% +%%% This program is distributed in the hope that it will be useful, +%%% but WITHOUT ANY WARRANTY; without even the implied warranty of +%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +%%% General Public License for more details. +%%% +%%% You should have received a copy of the GNU General Public License along +%%% with this program; if not, write to the Free Software Foundation, Inc., +%%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +%%% +%%%------------------------------------------------------------------- +-module(ejabberd_sm_riak). +-behaviour(ejabberd_sm). + +%% API +-export([init/0, set_session/1, delete_session/1, get_sessions/0, + get_sessions/1, get_sessions/2]). + +-include("ejabberd_sm.hrl"). +-include("logger.hrl"). + +%%%=================================================================== +%%% API +%%%=================================================================== +init() -> + clean_table(). + +set_session(Session) -> + ejabberd_riak:put(Session, session_schema(), + [{'2i', [{<<"us">>, Session#session.us}]}]). + +delete_session(Session) -> + ejabberd_riak:delete(session, Session#session.sid). + +get_sessions() -> + case ejabberd_riak:get(session, session_schema()) of + {ok, Ss} -> Ss; + {error, _} -> [] + end. + +get_sessions(LServer) -> + [S || S <- get_sessions(), element(2, S#session.us) == LServer]. + +get_sessions(U, S) -> + ejabberd_riak:get_by_index(session, session_schema(), <<"us">>, {U, S}). + +%%%=================================================================== +%%% Internal functions +%%%=================================================================== +session_schema() -> + {record_info(fields, session), #session{}}. + +clean_table() -> + %% TODO: not very efficient, rewrite using map-reduce or something + ?INFO_MSG("Cleaning Riak 'sm' table...", []), + lists:foreach( + fun(#session{sid = {_, Pid} = SID}) when node(Pid) == node() -> + ejabberd_riak:delete(session, SID); + (_) -> + ok + end, get_sessions()). diff --git a/test/ejabberd_SUITE_data/ejabberd.yml b/test/ejabberd_SUITE_data/ejabberd.yml index 3bfcfb235..24509737e 100644 --- a/test/ejabberd_SUITE_data/ejabberd.yml +++ b/test/ejabberd_SUITE_data/ejabberd.yml @@ -308,6 +308,7 @@ Welcome to this XMPP server." "riak.localhost": queue_type: ram auth_method: riak + sm_db_type: riak modules: mod_announce: db_type: riak