[STAFF-31] Add proper shutdown handling to Amqpx generic producers an…

…d consumers (#138)

* feat: implement a terminate for producers and consumers

which turns off the channel normally when reason is
`:normal`, `:shutdown`, or `{:shutdown, term}`

* chore: extract the opts for the genserver from the start_link param

* fixup! chore: extract the opts for the genserver from the start_link param

lib/amqp/gen/consumer.ex

@@ -24,11 +24,18 @@ defmodule Amqpx.Gen.Consumer do
   @callback handle_message_rejection(message :: any(), error :: any()) :: :ok | {:error, any()}
   @optional_callbacks handle_message_rejection: 2
 
+  @gen_server_opts [:name, :timeout, :debug, :spawn_opt, :hibernate_after]
+
+  @spec start_link(opts :: map()) :: GenServer.server()
   def start_link(opts) do
-    GenServer.start_link(__MODULE__, opts)
+    {gen_server_opts, opts} = Map.split(opts, @gen_server_opts)
+
+    GenServer.start_link(__MODULE__, opts, Map.to_list(gen_server_opts))
   end
 
   def init(opts) do
+    Process.flag(:trap_exit, true)
+
     state = struct(__MODULE__, opts)
     Process.send(self(), :setup, [])
     {:ok, state}
@@ -178,14 +185,28 @@ defmodule Amqpx.Gen.Consumer do
 
   def terminate(_, %__MODULE__{channel: nil}), do: nil
 
-  def terminate(_, %__MODULE__{channel: channel}) do
-    if Process.alive?(channel.pid) do
-      Channel.close(channel)
+  def terminate(reason, %__MODULE__{channel: channel}) do
+    case reason do
+      :normal -> close_channel(channel)
+      :shutdown -> close_channel(channel)
+      {:shutdown, _} -> close_channel(channel)
+      _ -> :ok
     end
   end
 
   # Private functions
 
+  defp close_channel(%{pid: pid} = channel) do
+    if Process.alive?(pid) do
+      Channel.close(channel)
+
+      receive do
+        {:DOWN, _, :process, ^pid, _reason} ->
+          :ok
+      end
+    end
+  end
+
   defp handle_message(
          message,
          %{delivery_tag: tag, redelivered: redelivered} = meta,

lib/amqp/gen/producer.ex

@@ -17,6 +17,8 @@ defmodule Amqpx.Gen.Producer do
 
   @default_backoff [base_ms: 10, max_ms: 5_000]
 
+  @gen_server_opts [:name, :timeout, :debug, :spawn_opt, :hibernate_after]
+
   defstruct [
     :channel,
     :publisher_confirms,
@@ -29,12 +31,20 @@ defmodule Amqpx.Gen.Producer do
 
   # Public API
 
+  @spec start_link(opts :: map()) :: GenServer.server()
   def start_link(opts) do
-    name = Map.get(opts, :name, __MODULE__)
-    GenServer.start_link(__MODULE__, opts, name: name)
+    gen_server_opts =
+      opts
+      |> Map.take(@gen_server_opts)
+      |> Map.to_list()
+      |> Keyword.put_new(:name, __MODULE__)
+
+    GenServer.start_link(__MODULE__, opts, gen_server_opts)
   end
 
   def init(opts) do
+    Process.flag(:trap_exit, true)
+
     case validate_configuration(opts) do
       {:ok, state} ->
         Process.send(self(), :setup, [])
@@ -123,9 +133,12 @@ defmodule Amqpx.Gen.Producer do
 
   def terminate(_, %__MODULE__{channel: nil}), do: nil
 
-  def terminate(_, %__MODULE__{channel: channel}) do
-    if Process.alive?(channel.pid) do
-      Channel.close(channel)
+  def terminate(reason, %__MODULE__{channel: channel}) do
+    case reason do
+      :normal -> close_channel(channel)
+      :shutdown -> close_channel(channel)
+      {:shutdown, _} -> close_channel(channel)
+      _ -> :ok
     end
   end
 
@@ -183,6 +196,18 @@ defmodule Amqpx.Gen.Producer do
   end
 
   # Private functions
+
+  defp close_channel(%{pid: pid} = channel) do
+    if Process.alive?(pid) do
+      Channel.close(channel)
+
+      receive do
+        {:DOWN, _, :process, ^pid, _reason} ->
+          :ok
+      end
+    end
+  end
+
   def do_publish_by(producer_name, exchange_name, routing_key, payload, options, attempt) do
     case GenServer.call(producer_name, {:publish, {exchange_name, routing_key, payload, options, attempt}}) do
       :ok ->

test/gen_test.exs

@@ -548,4 +548,48 @@ defmodule Amqpx.Test.AmqpxTest do
       end
     end
   end
+
+  test "the consumer should stop gracefully" do
+    tmp_ex = %{name: "tmp_ex", type: :topic, routing_keys: ["amqpx.tmp_ex"], opts: [durable: true]}
+
+    consumer_config = %{
+      queue: "tmp_q",
+      exchanges: [tmp_ex],
+      opts: [durable: true]
+    }
+
+    assert {:ok, test_chan} =
+             Amqpx.Gen.ConnectionManager
+             |> GenServer.call(:get_connection)
+             |> Amqpx.Channel.open(self())
+
+    assert :ok = Amqpx.Helper.setup_exchange(test_chan, tmp_ex)
+    assert :ok = Amqpx.Helper.setup_queue(test_chan, consumer_config)
+
+    test = self()
+
+    with_mock(Consumer1,
+      handle_message: fn _, _, s -> {:ok, s} end,
+      setup: fn channel ->
+        send(test, {:channel, channel})
+        {:ok, %{}}
+      end
+    ) do
+      {:ok, consumer_pid} = Amqpx.Gen.Consumer.start_link(%{handler_module: Consumer1})
+      Process.unlink(consumer_pid)
+
+      assert_receive {:channel, %{pid: channel_pid}}
+
+      Process.monitor(consumer_pid)
+      Process.monitor(channel_pid)
+
+      Process.exit(consumer_pid, {:shutdown, :die_gracefully})
+
+      assert_receive {:DOWN, _, :process, pid_1, _}
+      assert_receive {:DOWN, _, :process, pid_2, _}
+
+      assert pid_1 in [consumer_pid, channel_pid]
+      assert pid_2 in [consumer_pid, channel_pid]
+    end
+  end
 end