retort

defmodule Retort.Connection do
  alias Confex.Resolver

  @moduledoc """
  Caches the `AMQP.Connection` while it remains alive, so that multiple TCP connections aren't made to the AMQP broker,
  following best practices.

  Multiple channels should be opened on the same connection instead of multiple connections

      iex> {:ok, connection} = Retort.Connection.await
      iex> first_channel = AMQP.Channel.open(connection)
      iex> second_channel = AMQP.Channel.open(connection)
      iex> first_channel != second_channel
      true
  """

  # Constants

  @default_await_timeout :infinity
  @default_connection_timeout 5_000 # milliseconds

  use Connection

  require Logger

  # Struct

  defstruct connection: nil,
            waiter_from_by_monitor_reference: %{}

  # Types

  @typedoc """
  `connection` - the active connection to RabbitMQ.  `nil` when wait for initial connection or reconnect
  `waiter_from_by_monitor_reference` - tracks the list of `from` for calls to `await` by the `Process.monitor` `ref`
    used to monitor the waiter `pid` in `from`.  This is empty if there is an active `connection` when `await` is
    called; otherwise, it tracks all the processes that need to be notified (with `notify_waiters/1`) when a connection
    is finally established.  Entries are removed if the waiter dies, which is why its keyed by monitor reference.
  """
  @type t :: %__MODULE__{
               connection: %AMQP.Connection{} | nil,
               waiter_from_by_monitor_reference: %{
                 reference => GenServer.from
               }
             }

  # Functions

  ## Client API

  @doc """
  Waits for a connection, unlike `get/0`, which returns immediately if there is no connection.  Unlike get, which has
  default timeout of twice the `get_connection_timeout`, await's default timeout is `#{inspect @default_await_timeout}`.
  """
  @spec await :: {:ok, %AMQP.Connection{}} | no_return
  @spec await(timeout) :: {:ok, %AMQP.Connection{}} | no_return
  def await(timeout \\ @default_await_timeout), do: GenServer.call(__MODULE__, :await, timeout)

  @doc """
  Gets cached connection if it is still alive.  Otherwise, returns `{:error, :disconnected}`.  If you want to wait for
  a connection, use `await/1`

  # Gets the same `%AMQP.Connection{}` while it is alive

      iex> {:ok, original_connection} = Retort.Connection.get()
      iex> {:ok, original_connection} == Retort.Connection.get()
      true

  # Returns

  * `{:ok, %AMQP.Connection{}}` - connection is already established (or was established during the `get/1` timeout
    before the GenServer saw the get call)
  * `{:error, :disconnected}` - connection is not established.  Use `await/0` to wait until the a connection is
    available.
  """
  @spec get() :: {:ok, %AMQP.Connection{}} | {:error, :disconnected} | no_return
  def get do
    default_get_timeout()
    |> get()
  end

  @spec get(timeout) :: {:ok, %AMQP.Connection{}} | {:error, :disconnected} | no_return
  def get(timeout) do
    warn_timeout_inversion(timeout)

    GenServer.call(__MODULE__, :get, timeout)
  end

  @doc """
  The current backoff time (in milliseconds) for when `connect/2` fails to open an connection
  """
  @spec get_backoff() :: pos_integer
  def get_backoff do
    __MODULE__
    |> config_get([])
    |> Keyword.get(:backoff, get_connection_timeout())
  end

  @doc """
  The current connection timeout for `AMQP.Connection.open/1`
  """
  def get_connection_timeout do
    __MODULE__
    |> config_get([])
    |> Keyword.get(:connection_timeout, @default_connection_timeout)
  end

  @doc """
  The current URL to connect to RabbitMQ
  """
  @spec get_url :: String.t
  def get_url do
    :rabbitmq
    |> config_get()
    |> Keyword.fetch!(:url)
  end

  @doc """
  Updates the backoff for the next time one `connect/2` fails.
  """
  @spec put_backoff(pos_integer) :: :ok
  def put_backoff(backoff) when is_integer(backoff) and backoff > 0 do
    updated_env = __MODULE__
                  |> config_get([])
                  |> put_in([:backoff], backoff)
    Application.put_env(:retort, __MODULE__, updated_env)
  end

  @doc """
  Updates the connection timeout for `AMQP.Connection.open/1`
  """
  @spec put_connection_timeout(timeout) :: :ok
  def put_connection_timeout(timeout) when timeout == :infinity or is_integer(timeout) and timeout > 0 do
    updated_env = __MODULE__
                  |> config_get([])
                  |> put_in([:connection_timeout], timeout)
    Application.put_env(:retort, __MODULE__, updated_env)
  end

  @doc """
  Updates the URL to connect to RabbitMQ
  """
  @spec put_url(String.t) :: :ok
  def put_url(url) when is_binary(url) do
    updated_env = :rabbitmq
                  |> config_get()
                  |> put_in([:url], url)
    Application.put_env(:retort, :rabbitmq, updated_env)
  end

  @doc """
  Starts the connection cache.
  """
  @spec start_link(%AMQP.Connection{} | nil, GenServer.options) :: GenServer.on_start
  def start_link(connection, gen_server_opts \\ []), do: Connection.start_link(__MODULE__, connection, gen_server_opts)

  ## Connection Callbacks

  @doc """

  """
  @spec connect(:init | :reconnect, %__MODULE__{connection: nil}) ::
          {:ok, %__MODULE__{connection: %AMQP.Connection{}, waiter_from_by_monitor_reference: %{}}} |
          {:backoff, pos_integer, %__MODULE__{connection: nil, waiter_from_by_monitor_reference: map}}
  def connect(connect_reason, state) do
    case open() do
      {:ok, connection = %AMQP.Connection{pid: pid}} ->
        Logger.info fn ->
          ["Connected to RabbitMQ on ", inspect(connect_reason)]
        end

        Process.monitor(pid)

        state_without_waiters = %__MODULE__{state | connection: connection}
                                |> notify_waiters

        {:ok, state_without_waiters}
      {:error, error_reason} ->
        Logger.error fn ->
          ["Could not connect (", inspect(connect_reason), ") to RabbitMQ due to error (", inspect(error_reason), ")"]
        end

        {:backoff, get_backoff(), state}
    end
  end

  @doc """
  Initializes state to the given `conenction`.
  """
  @spec init(%AMQP.Connection{}) :: {:ok, %__MODULE__{connection: %AMQP.Connection{}}}
  @spec init(nil) :: {:connect, :init, %__MODULE__{connection: nil}}
  def init(connection) do
    Process.flag(:trap_exit, true)

    case connection do
      %AMQP.Connection{} -> {:ok, %__MODULE__{connection: connection}}
      nil -> {:connect, :init, %__MODULE__{connection: nil}}
    end
  end

  @spec handle_call(:await, any, t) :: {:reply, {:ok, %AMQP.Connection{}}, t} | {:noreply, t}
  @spec handle_call(:get, any, t) :: {:reply, {:error, :disconnected}, t} | {:reply, {:ok, %AMQP.Connection{}}, t}

  def handle_call(:await, _from, state = %__MODULE__{connection: connection = %AMQP.Connection{}}) do
    {:reply, {:ok, connection}, state}
  end

  def handle_call(:await, from, state = %__MODULE__{connection: nil}), do: {:noreply, put_waiter(state, from)}

  def handle_call(:get, _from, state = %__MODULE__{connection: nil}), do: {:reply, {:error, :disconnected}, state}

  def handle_call(:get, _from, state = %__MODULE__{}), do: {:reply, notification(state), state}

  @doc """
  Removes `%AMQP.Connection{}` from state when its `AMPQ.Connection.pid` goes down and tell Connection behaviour to
  reconnect.
  """
  @spec handle_info({:DOWN, reference, :process, pid, term} | term, t) :: {:connect, :reconnect, t} | {:noreply, t}

  def handle_info(
        {:DOWN, _, :process, pid, reason},
        state = %__MODULE__{
          connection: %AMQP.Connection{
            pid: pid
          }
        }
      ) do
    Logger.error "AMQP.Connection pid (#{inspect pid}) went down (#{inspect reason}).  Reconnecting..."

    {:connect, :reconnect, %__MODULE__{state | connection: nil}}
  end

  def handle_info(
        {:DOWN, monitor_reference, :process, pid, reason},
        state = %__MODULE__{
          waiter_from_by_monitor_reference: waiter_from_by_monitor_reference
        }
      ) do
    case Map.fetch(waiter_from_by_monitor_reference, monitor_reference) do
      {:ok, from} ->
        known_monitor_reference(from, monitor_reference, pid, reason, state)
      :error ->
        unknown_monitor_reference(monitor_reference, pid, state)
    end
  end

  # catchall clause
  def handle_info(info, state) do
    Logger.error "Got unexpected info #{inspect info}"

    {:noreply, state}
  end

  ## Private Functions

  @spec config_get(atom(), any()) :: any()
  defp config_get(key, default \\ nil) do
    :retort
    |> Application.get_env(key, default)
    |> Resolver.resolve!()
  end

  @spec default_get_timeout() :: non_neg_integer()
  defp default_get_timeout do
    # open/0 will block for `get_connection_timeout` and the GenServer can't respond to get calls during that time, so
    # get timeout *MUST* be greater than connection timeout to have any chance of succeeding when disconnected.
    get_connection_timeout() * 2
  end

  @spec known_monitor_reference(GenServer.from, reference, pid, any, t) :: {:noreply, t}
  defp known_monitor_reference({pid, _}, monitor_reference, pid, reason, state) do
    Logger.error "Process (#{inspect pid}) that was awaiting AMQP.Connection went down (#{inspect reason}).  " <>
                 "Removing the process from waiter_from_by_monitor_reference"

    new_state = update_in state.waiter_from_by_monitor_reference, fn waiter_from_by_monitor_reference ->
      Map.delete(waiter_from_by_monitor_reference, monitor_reference)
    end

    {:noreply, new_state}
  end

  defp known_monitor_reference({other_pid, _}, monitor_reference, pid, _reason, state) do
    Logger.error "Process (#{inspect pid}) went down, " <>
                 "but its monitor (#{inspect monitor_reference}) was registered " <>
                 "to another pid (#{inspect other_pid}) in waiter_from_monitor_reference"
    {:noreply, state}
  end

  defp notification(%__MODULE__{connection: connection = %AMQP.Connection{}}), do: {:ok, connection}

  @spec notify_waiter({reference, GenServer.from}, {:ok, %AMQP.Connection{}}) :: boolean
  defp notify_waiter({monitor_reference, from}, reply) do
    GenServer.reply(from, reply)
    Process.demonitor(monitor_reference, [:flush])
  end

  @spec notify_waiters(%__MODULE__{connection: %AMQP.Connection{}}) ::
          %__MODULE__{waiter_from_by_monitor_reference: %{}}
  defp notify_waiters(
         state = %__MODULE__{
           waiter_from_by_monitor_reference: waiter_from_by_monitor_reference
         }
       ) do
    reply = notification(state)
    Enum.each waiter_from_by_monitor_reference, &notify_waiter(&1, reply)

    %__MODULE__{state | waiter_from_by_monitor_reference: %{}}
  end

  defp open, do: AMQP.Connection.open("#{get_url()}?connection_timeout=#{get_connection_timeout()}")

  @spec put_waiter(t, GenServer.from) :: t
  defp put_waiter(state = %__MODULE__{}, from = {pid, _}) do
    monitor_reference = Process.monitor(pid)
    put_in state.waiter_from_by_monitor_reference[monitor_reference], from
  end

  @spec unknown_monitor_reference(reference, pid, t) :: {:noreply, t}
  defp unknown_monitor_reference(monitor_reference, pid, state) do
    Logger.error "#{inspect pid} went down, " <>
                 "but its monitor (#{inspect monitor_reference}) wasn't in waiter_from_by_monitor_reference.  " <>
                 "There is an unhandled link somwhere: use observer to find the link."
    {:noreply, state}
  end

  defp warn_timeout_inversion(timeout) do
    connection_timeout = get_connection_timeout()

    if timeout <= connection_timeout do
      Logger.warn(
        "#{__MODULE__}.get/1 timeout (#{timeout}) is less than or equal to " <>
        "connection timeout (#{connection_timeout}) and will fail when opening a new connection"
      )
    end
  end
end