Refactor mongoose_rabbit_worker

- Transform options directly into state instead of converting it
  into a proplist, and then into another map with different keys.
- Remove 'TODO Refactor' as we don't want to make the connection
  async. This could be an option added later if necessary,
  but currently it is consistent with other pools, e.g. rdbms.
- Add type specs for gen_server callbacks.
- Enable confirms after reconnection.
  So far, this would happen only on the first connection attempt.
- Use state and opts in more functions instead of additional data types.

Author: chrzaszcz

src/mongoose_rabbit_worker.erl

@@ -34,15 +34,20 @@
 
 -ignore_xref([start_link/0]).
 
--type state() :: #{amqp_client_opts := mongoose_amqp:network_params(),
-                   connection := pid(),
-                   channel := pid(),
-                   host_type := mongooseim:host_type_or_global(),
+-type state() :: #{host_type := mongooseim:host_type_or_global(),
                    pool_tag := atom(),
-                   confirms := boolean(),
-                   max_queue_len := non_neg_integer() | infinity}.
+                   opts := opts(),
+                   connection => pid(),
+                   channel => pid()}.
+
+-type opts() :: #{host := string(),
+                  port := inet:port_number(),
+                  username := binary(),
+                  password := binary(),
+                  virtual_host := binary(),
+                  confirms_enabled := boolean(),
+                  max_worker_queue_len := non_neg_integer() | infinity}.
 
--type worker_opts() :: state().
 -type publish_result() :: boolean() | timeout | {channel_exception, any(), any()}.
 
 %%%===================================================================
@@ -69,21 +74,28 @@ instrumentation(HostType, Tag) ->
 %%% gen_server callbacks
 %%%===================================================================
 
+-spec init(state()) -> {ok, state()}.
 init(Opts) ->
     process_flag(trap_exit, true),
-    %TODO: Refactor with handle_continue when OTP 21 is minimal supported version
-    do_init(Opts).
+    {ok, establish_rabbit_connection(Opts)}.
 
+-spec handle_call({amqp_call, mongoose_amqp:method()}, gen_server:from(), state()) ->
+          {reply, request_dropped | {ok | error | exit | throw, mongoose_amqp:method() | atom()},
+           state()}.
 handle_call(Req, From, State) ->
     maybe_handle_request(fun do_handle_call/3, [Req, From, State],
                          {reply, request_dropped, State}).
 
+-spec handle_cast({amqp_publish, mongoose_amqp:method(), mongoose_amqp:message()}, state()) ->
+          {noreply, state()}.
 handle_cast(Req, State) ->
-   maybe_handle_request(fun do_handle_cast/2, [Req, State], {noreply, State}).
+    maybe_handle_request(fun do_handle_cast/2, [Req, State], {noreply, State}).
 
+-spec handle_info(term(), state()) -> {noreply, state()}.
 handle_info(Req, State) ->
     maybe_handle_request(fun do_handle_info/2, [Req, State], {noreply, State}).
 
+-spec terminate(term(), state()) -> ok.
 terminate(_Reason, #{connection := Connection, channel := Channel,
                      host_type := HostType, pool_tag := PoolTag}) ->
     close_rabbit_connection(Connection, Channel, HostType, PoolTag),
@@ -93,55 +105,44 @@ terminate(_Reason, #{connection := Connection, channel := Channel,
 %%% Internal functions
 %%%===================================================================
 
-do_init(Opts) ->
-    HostType = proplists:get_value(host_type, Opts),
-    PoolTag = proplists:get_value(pool_tag, Opts),
-    AMQPClientOpts = proplists:get_value(amqp_client_opts, Opts),
-    {Connection, Channel} =
-        establish_rabbit_connection(AMQPClientOpts, HostType, PoolTag),
-    IsConfirmEnabled = maybe_enable_confirms(Channel, Opts),
-    MaxMsgQueueLen = proplists:get_value(max_queue_len, Opts),
-    {ok, #{host_type => HostType, amqp_client_opts => AMQPClientOpts,
-           connection => Connection, channel => Channel,
-           confirms => IsConfirmEnabled, max_queue_len => MaxMsgQueueLen,
-           pool_tag => PoolTag}}.
-
+-spec do_handle_call({amqp_call, mongoose_amqp:method()}, gen_server:from(), state()) ->
+          {reply, {ok | error | exit | throw, mongoose_amqp:method() | atom()}, state()}.
 do_handle_call({amqp_call, Method}, _From, State = #{channel := Channel}) ->
     try amqp_channel:call(Channel, Method) of
         Res ->
             {reply, {ok, Res}, State}
     catch
         Error:Reason ->
-            {FreshConn, FreshChann} = maybe_restart_rabbit_connection(State),
-            {reply, {Error, Reason}, State#{connection := FreshConn,
-                                            channel := FreshChann}}
+            {reply, {Error, Reason}, maybe_restart_rabbit_connection(State)}
     end.
 
+-spec do_handle_cast({amqp_publish, mongoose_amqp:method(), mongoose_amqp:message()}, state()) ->
+          {noreply, state()}.
 do_handle_cast({amqp_publish, Method, Payload}, State) ->
     handle_amqp_publish(Method, Payload, State).
 
+-spec do_handle_info(term(), state()) -> {noreply, state()}.
 do_handle_info(Req, State) ->
     ?UNEXPECTED_INFO(Req),
     {noreply, State}.
 
--spec handle_amqp_publish(Method :: mongoose_amqp:method(),
-                          Payload :: mongoose_amqp:message(),
-                          Opts :: worker_opts()) -> {noreply, worker_opts()}.
-handle_amqp_publish(Method, Payload, Opts = #{host_type := HostType,
-                                              pool_tag := PoolTag}) ->
-    PublishArgs = [Method, Payload, Opts],
-    Res = mongoose_instrument:span(wpool_rabbit_messages_published, #{host_type => HostType, pool_tag => PoolTag},
+-spec handle_amqp_publish(mongoose_amqp:method(), mongoose_amqp:message(), state()) ->
+          {noreply, state()}.
+handle_amqp_publish(Method, Payload, State = #{host_type := HostType, pool_tag := PoolTag}) ->
+    PublishArgs = [Method, Payload, State],
+    Res = mongoose_instrument:span(wpool_rabbit_messages_published,
+                                   #{host_type => HostType, pool_tag => PoolTag},
                                    fun publish_message_and_wait_for_confirm/3,
                                    PublishArgs,
                                    fun(PublishTime, Result) ->
-                                       handle_publish_result(PublishTime, Result, HostType, PoolTag, PublishArgs)
+                                       handle_publish_result(PublishTime, Result, HostType,
+                                                             PoolTag, PublishArgs)
                                    end),
     case Res of
         {channel_exception, _, _} ->
-            {FreshConn, FreshChann} = maybe_restart_rabbit_connection(Opts),
-            {noreply, Opts#{connection := FreshConn, channel := FreshChann}};
+            {noreply, maybe_restart_rabbit_connection(State)};
         _ ->
-            {noreply, Opts}
+            {noreply, State}
     end.
 
 -spec handle_publish_result(integer(), publish_result(), mongooseim:host_type_or_global(),
@@ -150,103 +151,97 @@ handle_amqp_publish(Method, Payload, Opts = #{host_type := HostType,
 handle_publish_result(PublishTime, true, HostType, PoolTag, [Method, Payload, Opts]) ->
     ?LOG_DEBUG(#{what => rabbit_message_sent, host_type => HostType, tag => PoolTag,
                  method => Method, payload => Payload, opts => Opts}),
-    #{count => 1, time => PublishTime, size => byte_size(term_to_binary(Payload)), payload => Payload};
+    #{count => 1, time => PublishTime, size => byte_size(term_to_binary(Payload)),
+      payload => Payload};
 handle_publish_result(_PublishTime, false, HostType, PoolTag, [Method, Payload, Opts]) ->
-    ?LOG_WARNING(#{what => rabbit_message_sent_failed, reason => negative_ack, host_type => HostType, tag => PoolTag,
-                   method => Method, payload => Payload, opts => Opts}),
+    ?LOG_WARNING(#{what => rabbit_message_sent_failed, reason => negative_ack,
+                   host_type => HostType, tag => PoolTag, method => Method, payload => Payload,
+                   opts => Opts}),
     #{failed => 1, payload => Payload};
-handle_publish_result(_PublishTime, {channel_exception, Error, Reason}, HostType, PoolTag, [Method, Payload, Opts]) ->
+handle_publish_result(_PublishTime, {channel_exception, Error, Reason}, HostType, PoolTag,
+                      [Method, Payload, Opts]) ->
     ?LOG_ERROR(#{what => rabbit_message_sent_failed,
                  class => Error, reason => Reason, host_type => HostType, tag => PoolTag,
                  method => Method, payload => Payload, opts => Opts}),
     #{failed => 1, payload => Payload};
 handle_publish_result(_PublishTime, timeout, HostType, PoolTag, [Method, Payload, Opts]) ->
-    ?LOG_ERROR(#{what => rabbit_message_sent_failed, reason => timeout, host_type => HostType, tag => PoolTag,
-                 method => Method, payload => Payload, opts => Opts}),
+    ?LOG_ERROR(#{what => rabbit_message_sent_failed, reason => timeout, host_type => HostType,
+                 tag => PoolTag, method => Method, payload => Payload, opts => Opts}),
     #{timeout => 1, payload => Payload}.
 
 -spec publish_message_and_wait_for_confirm(Method :: mongoose_amqp:method(),
-                                           Payload :: mongoose_amqp:message(),
-                                           worker_opts()) ->
+                                           Payload :: mongoose_amqp:message(), state()) ->
     publish_result().
-publish_message_and_wait_for_confirm(Method, Payload,
-                                     #{channel := Channel,
-                                       confirms := IsConfirmEnabled}) ->
+publish_message_and_wait_for_confirm(Method, Payload, #{channel := Channel, opts := Opts}) ->
     try amqp_channel:call(Channel, Method, Payload) of
         _Res ->
-            maybe_wait_for_confirms(Channel, IsConfirmEnabled)
+            maybe_wait_for_confirms(Channel, Opts)
     catch
         Error:Reason -> {channel_exception, Error, Reason}
     end.
 
--spec maybe_wait_for_confirms(Channel :: pid(), boolean()) ->
-    boolean() | timeout.
-maybe_wait_for_confirms(Channel, true) ->
+-spec maybe_wait_for_confirms(Channel :: pid(), opts()) -> boolean() | timeout.
+maybe_wait_for_confirms(Channel, #{confirms_enabled := true}) ->
     amqp_channel:wait_for_confirms(Channel);
-maybe_wait_for_confirms(_, _) -> true.
+maybe_wait_for_confirms(_, _) ->
+    true.
 
--spec maybe_restart_rabbit_connection(worker_opts()) -> {pid(), pid()}.
-maybe_restart_rabbit_connection(#{connection := Conn, host_type := HostType,
-                                  pool_tag := PoolTag,
-                                  amqp_client_opts := AMQPOpts}) ->
+-spec maybe_restart_rabbit_connection(state()) -> state().
+maybe_restart_rabbit_connection(#{connection := Conn} = State) ->
     case is_process_alive(Conn) of
         true ->
             {ok, Channel} = amqp_connection:open_channel(Conn),
-            {Conn, Channel};
+            State#{channel := Channel};
         false ->
-            establish_rabbit_connection(AMQPOpts, HostType, PoolTag)
+            establish_rabbit_connection(State)
     end.
 
--spec establish_rabbit_connection(Opts :: mongoose_amqp:network_params(),
-                                  HostType :: mongooseim:host_type_or_global(), PoolTag :: atom())
-    -> {pid(), pid()}.
-establish_rabbit_connection(AMQPOpts, HostType, PoolTag) ->
-    case amqp_connection:start(AMQPOpts) of
+-spec establish_rabbit_connection(state()) -> state().
+establish_rabbit_connection(State) ->
+    #{opts := Opts, host_type := HostType, pool_tag := PoolTag} = State,
+    case amqp_connection:start(mongoose_amqp:network_params(Opts)) of
         {ok, Connection} ->
-            mongoose_instrument:execute(wpool_rabbit_connections, #{host_type => HostType, pool_tag => PoolTag},
+            mongoose_instrument:execute(wpool_rabbit_connections,
+                                        #{host_type => HostType, pool_tag => PoolTag},
                                         #{active => 1, opened => 1}),
             {ok, Channel} = amqp_connection:open_channel(Connection),
+            maybe_enable_confirms(Channel, Opts),
             ?LOG_DEBUG(#{what => rabbit_connection_established,
-                         host_type => HostType, pool_tag => PoolTag, opts => AMQPOpts}),
-            {Connection, Channel};
+                         host_type => HostType, pool_tag => PoolTag, opts => Opts}),
+            State#{connection => Connection, channel => Channel};
         {error, Error} ->
-            mongoose_instrument:execute(wpool_rabbit_connections, #{host_type => HostType, pool_tag => PoolTag},
+            mongoose_instrument:execute(wpool_rabbit_connections,
+                                        #{host_type => HostType, pool_tag => PoolTag},
                                         #{failed => 1}),
             ?LOG_ERROR(#{what => rabbit_connection_failed, reason => Error,
-                         host_type => HostType, pool_tag => PoolTag, opts => AMQPOpts}),
+                         host_type => HostType, pool_tag => PoolTag, opts => Opts}),
             exit("connection to a Rabbit server failed")
     end.
 
 -spec close_rabbit_connection(Connection :: pid(), Channel :: pid(),
                               HostType :: mongooseim:host_type_or_global(), PoolTag :: atom()) ->
     ok | no_return().
 close_rabbit_connection(Connection, Channel, HostType, PoolTag) ->
-    mongoose_instrument:execute(wpool_rabbit_connections, #{host_type => HostType, pool_tag => PoolTag},
+    mongoose_instrument:execute(wpool_rabbit_connections,
+                                #{host_type => HostType, pool_tag => PoolTag},
                                 #{active => -1, closed => 1}),
     try amqp_channel:close(Channel)
     catch
         _Error:_Reason -> already_closed
     end,
     amqp_connection:close(Connection).
 
--spec maybe_enable_confirms(Channel :: pid(), proplists:proplist()) ->
-    boolean() | no_return().
-maybe_enable_confirms(Channel, Opts) ->
-    case proplists:get_value(confirms, Opts) of
-        true ->
-            ConfirmCallRes = mongoose_amqp:confirm_select_ok(),
-            ConfirmCallRes =
-                amqp_channel:call(Channel, mongoose_amqp:confirm_select()),
-            true;
-        false ->
-            false
-    end.
+-spec maybe_enable_confirms(Channel :: pid(), opts()) -> ok.
+maybe_enable_confirms(Channel, #{confirms_enabled := true}) ->
+    ConfirmCallRes = mongoose_amqp:confirm_select_ok(),
+    ConfirmCallRes = amqp_channel:call(Channel, mongoose_amqp:confirm_select()),
+    ok;
+maybe_enable_confirms(_Channel, #{}) ->
+    ok.
 
--spec maybe_handle_request(Callback :: function(), Args :: [term()],
-                           Reply :: term()) -> term().
+-spec maybe_handle_request(Callback :: function(), Args :: [term()], Reply :: term()) -> term().
 maybe_handle_request(Callback, Args, Reply) ->
-    State = lists:last(Args),
-    Limit = maps:get(max_queue_len, State),
+    #{opts := #{max_worker_queue_len := Limit}} = lists:last(Args),
     case is_msq_queue_max_limit_reached(Limit) of
         false ->
             apply(Callback, Args);
@@ -257,8 +252,7 @@ maybe_handle_request(Callback, Args, Reply) ->
             Reply
     end.
 
--spec is_msq_queue_max_limit_reached(Limit :: infinity | non_neg_integer()) ->
-    boolean().
+-spec is_msq_queue_max_limit_reached(Limit :: infinity | non_neg_integer()) -> boolean().
 is_msq_queue_max_limit_reached(infinity) -> false;
 is_msq_queue_max_limit_reached(Limit) ->
     case process_info(self(), message_queue_len) of

src/wpool/mongoose_wpool_rabbit.erl

@@ -12,14 +12,11 @@ init() ->
 -spec start(mongooseim:host_type_or_global(), mongoose_wpool:tag(),
             mongoose_wpool:pool_opts(), mongoose_wpool:conn_opts()) -> {ok, pid()} | {error, any()}.
 start(HostType, Tag, WpoolOptsIn, ConnOpts) ->
-    #{confirms_enabled := Confirms, max_worker_queue_len := MaxQueueLen} = ConnOpts,
     PoolName = mongoose_wpool:make_pool_name(rabbit, HostType, Tag),
     Worker = {mongoose_rabbit_worker,
-              [{amqp_client_opts, mongoose_amqp:network_params(ConnOpts)},
-               {host_type, HostType},
-               {pool_tag, Tag},
-               {confirms, Confirms},
-               {max_queue_len, MaxQueueLen}]},
+              #{host_type => HostType,
+                pool_tag => Tag,
+                opts => ConnOpts}},
     WpoolOpts = [{worker, Worker}, {pool_sup_shutdown, timer:seconds(5)} | WpoolOptsIn],
     mongoose_wpool:start_sup_pool(rabbit, PoolName, WpoolOpts).
 

test/mongoose_rabbit_worker_SUITE.erl

@@ -54,16 +54,15 @@ end_per_suite(Config) ->
     Config.
 
 init_per_testcase(_Case, Config) ->
-    WorkerOpts = [{host_type, ?HOST_TYPE},
-                  {pool_tag, test_tag},
-                  {amqp_client_opts, []},
-                  {confirms, false},
-                  {max_queue_len, ?MAX_QUEUE_LEN}],
     mongoose_instrument:start_link(),
     mongoose_instrument:set_up(mongoose_wpool_rabbit:instrumentation(?HOST_TYPE, test_tag)),
+    WorkerOpts = #{host_type => ?HOST_TYPE, pool_tag => test_tag, opts => conn_opts()},
     {ok, WorkerPid} = gen_server:start(mongoose_rabbit_worker, WorkerOpts, []),
     Config ++ [{worker_pid, WorkerPid}].
 
+conn_opts() ->
+    config_parser_helper:default_config([outgoing_pools, rabbit, test_tag, conn_opts]).
+
 end_per_testcase(_Case, Config) ->
     mongoose_instrument:tear_down(mongoose_wpool_rabbit:instrumentation(?HOST_TYPE, test_tag)),
     exit(proplists:get_value(worker_pid, Config), kill),