Merge pull request #3778 from esl/c2s/amp-big

Fixes AMP big suite

big_tests/default.spec

@@ -17,7 +17,7 @@
 {suites, "tests", acc_e2e_SUITE}.
 {suites, "tests", accounts_SUITE}.
 {suites, "tests", adhoc_SUITE}.
-% {suites, "tests", amp_big_SUITE}.
+{suites, "tests", amp_big_SUITE}.
 % {suites, "tests", anonymous_SUITE}.
 {suites, "tests", auth_methods_for_c2s_SUITE}.
 {suites, "tests", bosh_SUITE}.

big_tests/dynamic_domains.spec

@@ -17,7 +17,7 @@
 
 {suites, "tests", adhoc_SUITE}.
 
-% {suites, "tests", amp_big_SUITE}.
+{suites, "tests", amp_big_SUITE}.
 
 % {suites, "tests", anonymous_SUITE}.
 

big_tests/tests/amp_big_SUITE.erl

@@ -124,8 +124,8 @@ amp_test_helper_code() ->
     "-module(amp_test_helper).\n"
     "-compile([export_all, nowarn_export_all]).\n"
     "setup_meck() ->\n"
-    "  meck:expect(ejabberd_socket, send, fun ejabberd_socket_send/2).\n"
-    "ejabberd_socket_send(Socket, Data) ->\n"
+    "  meck:expect(ranch_tcp, send, fun ranch_tcp_send/2).\n"
+    "ranch_tcp_send(Socket, Data) ->\n"
     "  case catch binary:match(Data, <<\"Recipient connection breaks\">>) of\n"
     "    {N, _} when is_integer(N) -> {error, simulated};\n"
     "    _ -> meck:passthrough([Socket, Data])\n"
@@ -149,7 +149,7 @@ init_per_group(GroupName, Config) ->
     save_offline_status(GroupName, Config1).
 
 setup_meck(suite) ->
-    ok = rpc(mim(), meck, new, [ejabberd_socket, [passthrough, no_link]]),
+    ok = rpc(mim(), meck, new, [ranch_tcp, [passthrough, no_link]]),
     ok = rpc(mim(), amp_test_helper, setup_meck, []);
 setup_meck(mam_failure) ->
     ok = rpc(mim(), meck, expect, [mod_mam_rdbms_arch, archive_message, 3, {error, simulated}]);

src/c2s/mongoose_c2s.erl

@@ -803,12 +803,14 @@ sm_unset_reason(_) ->
     error.
 
 %% @doc This is the termination point - from here stanza is sent to the user
--spec send_element(c2s_data(), exml:element(), mongoose_acc:t()) -> maybe_ok().
-send_element(StateData = #c2s_data{host_type = <<>>}, El, _) ->
-    send_xml(StateData, El);
+-spec send_element(c2s_data(), exml:element(), mongoose_acc:t()) -> mongoose_acc:t().
+send_element(StateData = #c2s_data{host_type = <<>>}, El, Acc) ->
+    send_xml(StateData, El),
+    Acc;
 send_element(StateData = #c2s_data{host_type = HostType}, El, Acc) ->
-    mongoose_hooks:xmpp_send_element(HostType, Acc, El),
-    send_xml(StateData, El).
+    Res = send_xml(StateData, El),
+    Acc1 = mongoose_acc:set(c2s, send_result, Res, Acc),
+    mongoose_hooks:xmpp_send_element(HostType, Acc1, El).
 
 -spec send_xml(c2s_data(), exml_stream:element() | [exml_stream:element()]) -> maybe_ok().
 send_xml(StateData, Xml) ->

src/ejabberd_sm_redis.erl

@@ -25,7 +25,9 @@
          total_count/0,
          unique_count/0]).
 
--ignore_xref([maybe_initial_cleanup/2]).
+-export([total_count_internal/0]).
+
+-ignore_xref([maybe_initial_cleanup/2, total_count_internal/0]).
 
 -spec init(map()) -> any().
 init(_Opts) ->
@@ -156,8 +158,15 @@ maybe_initial_cleanup(Node, Initial) ->
 
 -spec total_count() -> integer().
 total_count() ->
-    {Counts, _} = rpc:multicall(supervisor, count_children, [ejabberd_c2s_sup]),
-    lists:sum([proplists:get_value(active, Count, 0) || Count <- Counts, is_list(Count)]).
+    {Counts, _} = rpc:multicall(?MODULE, total_count_internal, []),
+    lists:sum(Counts).
+
+-spec total_count_internal() -> integer().
+
+total_count_internal() ->
+    Children = supervisor:which_children(mongoose_listener_sup),
+    Listeners = [Ref || {Ref, _, _, [mongoose_c2s_listener]} <- Children],
+    lists:sum([maps:get(active_connections, ranch:info(Ref)) || Ref <- Listeners]).
 
 -spec unique_count() -> integer().
 unique_count() ->

src/mod_amp.erl

@@ -9,13 +9,17 @@
 -behaviour(mongoose_module_metrics).
 -xep([{xep, 79}, {version, "1.2"}, {comment, "partially implemented."}]).
 -export([start/2, stop/1, supported_features/0]).
--export([run_initial_check/2,
+-export([run_initial_check/3,
          check_packet/2,
          disco_local_features/1,
-         c2s_stream_features/3
+         c2s_stream_features/3,
+         xmpp_send_element/2
         ]).
 
--ignore_xref([c2s_stream_features/3, disco_local_features/1, run_initial_check/2]).
+-ignore_xref([c2s_stream_features/3,
+              disco_local_features/1,
+              run_initial_check/2,
+              xmpp_send_element/2]).
 
 -include("amp.hrl").
 -include("mongoose.hrl").
@@ -29,32 +33,40 @@
 
 -spec start(mongooseim:host_type(), gen_mod:module_opts()) -> ok.
 start(HostType, _Opts) ->
-    ejabberd_hooks:add(hooks(HostType)).
+    ejabberd_hooks:add(hooks(HostType)),
+    gen_hook:add_handlers(c2s_hooks(HostType)).
 
 -spec stop(mongooseim:host_type()) -> ok.
 stop(HostType) ->
-    ejabberd_hooks:delete(hooks(HostType)).
+    ejabberd_hooks:delete(hooks(HostType)),
+    gen_hook:delete_handlers(c2s_hooks(HostType)).
 
 -spec supported_features() -> [atom()].
 supported_features() -> [dynamic_domains].
 
+-spec hooks(mongooseim:host_type()) -> [ejabberd_hooks:hook()].
 hooks(HostType) ->
     [{c2s_stream_features, HostType, ?MODULE, c2s_stream_features, 50},
      {disco_local_features, HostType, ?MODULE, disco_local_features, 99},
-     {c2s_preprocessing_hook, HostType, ?MODULE, run_initial_check, 10},
      {amp_verify_support, HostType, ?AMP_RESOLVER, verify_support, 10},
      {amp_check_condition, HostType, ?AMP_RESOLVER, check_condition, 10},
-     {amp_determine_strategy, HostType, ?AMP_STRATEGY, determine_strategy, 10}].
+     {amp_determine_strategy, HostType, ?AMP_STRATEGY, determine_strategy, 10},
+     {xmpp_send_element, HostType, ?MODULE, xmpp_send_element, 10}].
+
+-spec c2s_hooks(mongooseim:host_type()) -> gen_hook:hook_list(mongoose_c2s_hooks:hook_fn()).
+c2s_hooks(HostType) ->
+    [{c2s_preprocessing_hook, HostType, fun ?MODULE:run_initial_check/3, #{}, 10}].
 
 %% API
 
--spec run_initial_check(mongoose_acc:t(), ejabberd_c2s:state()) -> mongoose_acc:t().
-run_initial_check(#{result := drop} = Acc, _C2SState) ->
-    Acc;
-run_initial_check(Acc, _C2SState) ->
+-spec run_initial_check(mongoose_acc:t(), mongoose_c2s:hook_params(), gen_hook:extra()) ->
+    mongoose_c2s_hooks:hook_result().
+run_initial_check(#{result := drop} = Acc, _Params, _Extra) ->
+    {ok, Acc};
+run_initial_check(Acc, _Params, _Extra) ->
     case mongoose_acc:stanza_name(Acc) of
-        <<"message">> -> run_initial_check(Acc);
-        _ -> Acc
+        <<"message">> -> {ok, run_initial_check(Acc)};
+        _ -> {ok, Acc}
     end.
 
 -spec check_packet(mongoose_acc:t(), amp_event()) -> mongoose_acc:t().
@@ -236,3 +248,11 @@ find(Pred, [H|T]) ->
         true -> {found, H};
         false -> find(Pred, T)
     end.
+
+-spec xmpp_send_element(mongoose_acc:t(), exml:element()) -> mongoose_acc:t().
+xmpp_send_element(Acc, _El) ->
+    Event = case mongoose_acc:get(c2s, send_result, Acc) of
+                ok -> delivered;
+                _ -> delivery_failed
+            end,
+    check_packet(Acc, Event).