diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl index 97b171f52..db9e52bee 100644 --- a/src/ejabberd_c2s.erl +++ b/src/ejabberd_c2s.erl @@ -712,7 +712,7 @@ wait_for_auth({xmlstreamelement, El}, StateData) -> conn = Conn, auth_module = AuthModule, - pres_f = + pres_f = (?SETS):from_list(Fs1), pres_t = (?SETS):from_list(Ts1), @@ -2381,11 +2381,7 @@ presence_broadcast_to_trusted(StateData, From, Trusted, JIDSet, Packet) -> %% Send presence when connecting presence_broadcast_first(From, StateData, Packet) -> - JIDsProbe = - ?SETS:fold( - fun(JID, L) -> [JID | L] end, - [], - StateData#state.pres_t), + JIDsProbe = ?SETS:to_list(StateData#state.pres_t), PacketProbe = #xmlel{name = <<"presence">>, attrs = [{<<"type">>,<<"probe">>}], children = []}, JIDs2Probe = format_and_check_privacy(From, StateData, Packet, JIDsProbe, out), Server = StateData#state.server, @@ -3434,13 +3430,14 @@ pack(S = #state{pres_a = A, pres_f = F, {NewA, Pack2} = pack_jid_set(A, gb_trees:empty()), {NewF, Pack3} = pack_jid_set(F, Pack2), {NewT, _Pack4} = pack_jid_set(T, Pack3), - S#state{pres_a = NewA, pres_f = NewF, - pres_t = NewT}. + {SetF, SetT} = ?SETS:pack_sets(NewF, NewT), + S#state{pres_a = ?SETS:from_list(NewA), pres_f = SetF, + pres_t = SetT}. pack_jid_set(Set, Pack) -> Jids = (?SETS):to_list(Set), {PackedJids, NewPack} = pack_jids(Jids, Pack, []), - {(?SETS):from_list(PackedJids), NewPack}. + {PackedJids, NewPack}. pack_jids([], Pack, Acc) -> {Acc, Pack}; pack_jids([{U, S, R} = Jid | Jids], Pack, Acc) -> diff --git a/src/ejabberd_c2s.hrl b/src/ejabberd_c2s.hrl index 46e3e97b5..b8120b9ee 100644 --- a/src/ejabberd_c2s.hrl +++ b/src/ejabberd_c2s.hrl @@ -25,7 +25,8 @@ -endif. --define(SETS, gb_sets). +%-define(SETS, gb_sets). +-define(SETS, ejabberd_sets). -define(DICT, dict). diff --git a/src/ejabberd_sets.erl b/src/ejabberd_sets.erl new file mode 100644 index 000000000..90f2ec280 --- /dev/null +++ b/src/ejabberd_sets.erl @@ -0,0 +1,81 @@ +-module(ejabberd_sets). + +%% Sets with adaptable storage, use less memory if the set size isn't big. +%% +%% +%% -Use a simple lists, change to gb_sets only if size grows above thereshold +%% -Function to build two sets at once, from two lists, trying to share as much as +%% possible. This avoid the intersection to be repeated on both sets structure. + +%%SET interface, same than gb_sets +-export([ + new/0, + from_list/1, + to_list/1, + is_element/2, + add_element/2, + del_element/2, + foldl/3, + size/1]). + +%% new method +-export([ + pack_sets/2 + ]). + +-define(THRESHOLD, 100). + + +new() -> + []. +from_list(L) when length(L) < ?THRESHOLD -> + L; +from_list(L) -> + gb_sets:from_list(L). + +pack_sets(L1, L2) -> + Shared = lists:filter(fun(I) -> lists:member(I, L2) end, L1), + SharedSet = from_list(Shared), + %% add_element already verify that the elements doesn't exists, so no need to + %% filter L1 and L2 before. + R1 = lists:foldl( fun add_element/2, SharedSet, L1), + R2 = lists:foldl( fun add_element/2, SharedSet, L2), + {R1, R2}. + + + +to_list(S) when is_list(S) -> S; +to_list(S) -> gb_sets:to_list(S). + +is_element(El, S) when is_list(S) -> lists:member(El, S); +is_element(El, S) -> gb_sets:is_element(El, S). + +add_element(El, S) when is_list(S) -> + case lists:member(El, S) of + true -> + S; + false -> + fix_storage([El | S]) + end; +add_element(El, S) -> + gb_sets:add_element(El, S). + + + +size(S) when is_list(S) -> length(S); +size(S) -> gb_sets:size(S). + +foldl(Fun, Init, S) when is_list(S)-> lists:foldl(Fun, Init, S); +foldl(Fun, Init, S) -> gb_sets:foldl(Fun, Init, S). + +del_element(El, S) when is_list(S) -> + lists:remove(El, S); +del_element(El, S) -> + fix_storage(gb_sets:del_element(El,S)). + + + +fix_storage(S) when is_list(S), length(S) > ?THRESHOLD -> + gb_sets:from_list(S); +fix_storage(S) -> S. + %%For now we don't shrink to lists again, not neccesarly