db_connection

defmodule DBConnection.Holder do
  @moduledoc false
  require Record

  @queue true
  @timeout 15000
  @time_unit 1000

  Record.defrecord(:conn, [:connection, :module, :state, :lock, deadline: nil, status: :ok])
  Record.defrecord(:pool_ref, [:pool, :reference, :deadline, :holder, :lock])

  @type t :: :ets.tid()

  ## Holder API

  @spec new(pid, reference, module, term) :: t
  def new(pool, ref, mod, state) do
    # Insert before setting heir so that pool can't receive empty table
    holder = :ets.new(__MODULE__, [:public, :ordered_set])

    conn = conn(connection: self(), module: mod, state: state)
    true = :ets.insert_new(holder, conn)

    :ets.setopts(holder, {:heir, pool, ref})
    holder
  end

  @spec update(pid, reference, module, term) :: {:ok, t} | :error
  def update(pool, ref, mod, state) do
    holder = new(pool, ref, mod, state)
    now = System.monotonic_time(@time_unit)

    try do
      :ets.give_away(holder, pool, {:checkin, ref, now})
      {:ok, holder}
    rescue
      ArgumentError -> :error
    end
  end

  @spec delete(t) :: {module, term}
  def delete(holder) do
    [conn(module: module, state: state)] = :ets.lookup(holder, :conn)
    :ets.delete(holder)
    {module, state}
  end

  ## Pool API (invoked by caller)

  @callback checkout(pool :: GenServer.server(), opts :: Keyword.t()) ::
              {:ok, pool_ref :: any, module, state :: any} | {:error, Exception.t()}
  def checkout(pool, opts) do
    caller = Keyword.get(opts, :caller, self())
    callers = [caller | Process.get(:"$callers") || []]

    queue? = Keyword.get(opts, :queue, @queue)
    now = System.monotonic_time(@time_unit)
    timeout = abs_timeout(now, opts)

    case checkout(pool, callers, queue?, now, timeout) do
      {:ok, _, _, _} = ok ->
        ok

      {:error, _} = error ->
        error

      {:redirect, proxy} ->
        checkout(proxy, opts)

      {:exit, reason} ->
        exit({reason, {__MODULE__, :checkout, [pool, opts]}})
    end
  end

  @spec checkin(pool_ref :: any) :: :ok
  def checkin(pool_ref) do
    now = System.monotonic_time(@time_unit)
    # Note we may call checkin after a disconnect/stop. For this reason, we choose
    # to not change the status on checkin but strictly speaking nobody can access
    # the holder after disconnect/stop unless they store a copy of %DBConnection{}.
    # Note status can't be :aborted as aborted is always reverted at the of a transaction.
    done(pool_ref, [{conn(:lock) + 1, nil}], :checkin, now)
  end

  @spec disconnect(pool_ref :: any, err :: Exception.t()) :: :ok
  def disconnect(pool_ref, err) do
    done(pool_ref, [{conn(:status) + 1, :error}], :disconnect, err)
  end

  @spec stop(pool_ref :: any, err :: Exception.t()) :: :ok
  def stop(pool_ref, err) do
    done(pool_ref, [{conn(:status) + 1, :error}], :stop, err)
  end

  @spec handle(pool_ref :: any, fun :: atom, args :: [term], Keyword.t()) :: tuple
  def handle(pool_ref, fun, args, opts) do
    handle_or_cleanup(:handle, pool_ref, fun, args, opts)
  end

  @spec cleanup(pool_ref :: any, fun :: atom, args :: [term], Keyword.t()) :: tuple
  def cleanup(pool_ref, fun, args, opts) do
    handle_or_cleanup(:cleanup, pool_ref, fun, args, opts)
  end

  defp handle_or_cleanup(type, pool_ref, fun, args, opts) do
    pool_ref(holder: holder, lock: lock) = pool_ref

    try do
      :ets.lookup(holder, :conn)
    rescue
      ArgumentError ->
        msg = "connection is closed because of an error, disconnect or timeout"
        {:disconnect, DBConnection.ConnectionError.exception(msg), _state = :unused}
    else
      [conn(lock: conn_lock)] when conn_lock != lock ->
        raise "an outdated connection has been given to DBConnection on #{fun}/#{length(args) + 2}"

      [conn(status: :error)] ->
        msg = "connection is closed because of an error, disconnect or timeout"
        {:disconnect, DBConnection.ConnectionError.exception(msg), _state = :unused}

      [conn(status: :aborted)] when type != :cleanup ->
        msg = "transaction rolling back"
        {:disconnect, DBConnection.ConnectionError.exception(msg), _state = :unused}

      [conn(module: module, state: state)] ->
        holder_apply(holder, module, fun, args ++ [opts, state])
    end
  end

  ## Pool state helpers API (invoked by callers)

  @spec put_state(pool_ref :: any, term) :: :ok
  def put_state(pool_ref(holder: sink_holder), state) do
    :ets.update_element(sink_holder, :conn, [{conn(:state) + 1, state}])
    :ok
  end

  @spec status?(pool_ref :: any, :ok | :aborted) :: boolean()
  def status?(pool_ref(holder: holder), status) do
    try do
      :ets.lookup_element(holder, :conn, conn(:status) + 1) == status
    rescue
      ArgumentError -> false
    end
  end

  @spec put_status(pool_ref :: any, :ok | :aborted) :: boolean()
  def put_status(pool_ref(holder: holder), status) do
    try do
      :ets.update_element(holder, :conn, [{conn(:status) + 1, status}])
    rescue
      ArgumentError -> false
    end
  end

  ## Pool callbacks (invoked by pools)

  @spec reply_redirect({pid, reference}, GenServer.server()) :: :ok
  def reply_redirect(from, redirect) do
    GenServer.reply(from, {:redirect, redirect})
    :ok
  end

  @spec reply_error({pid, reference}, Exception.t()) :: :ok
  def reply_error(from, exception) do
    GenServer.reply(from, {:error, exception})
    :ok
  end

  @spec handle_checkout(t, {pid, reference}, reference) :: boolean
  def handle_checkout(holder, {pid, mref}, ref) do
    :ets.give_away(holder, pid, {mref, ref})
  rescue
    ArgumentError ->
      if Process.alive?(pid) do
        raise ArgumentError, no_holder(holder)
      else
        false
      end
  end

  @spec handle_deadline(t, reference) :: boolean
  def handle_deadline(holder, deadline) do
    :ets.lookup_element(holder, :conn, conn(:deadline) + 1)
  rescue
    ArgumentError -> false
  else
    ^deadline -> true
    _ -> false
  end

  @spec handle_ping(t) :: boolean
  def handle_ping(holder) do
    :ets.lookup(holder, :conn)
  rescue
    ArgumentError ->
      raise ArgumentError, no_holder(holder)
  else
    [conn(connection: conn, state: state)] ->
      DBConnection.Connection.ping({conn, holder}, state)
      :ets.delete(holder)
  end

  @spec handle_disconnect(t, Exception.t()) :: :ok
  def handle_disconnect(holder, err) do
    handle_done(holder, &DBConnection.Connection.disconnect/3, err)
  end

  @spec handle_stop(t, term) :: :ok
  def handle_stop(holder, err) do
    handle_done(holder, &DBConnection.Connection.stop/3, err)
  end

  ## Private

  defp checkout(pool, callers, queue?, start, timeout) do
    case GenServer.whereis(pool) do
      pid when node(pid) == node() ->
        checkout_call(pid, callers, queue?, start, timeout)

      pid when node(pid) != node() ->
        {:exit, {:badnode, node(pid)}}

      {_name, node} ->
        {:exit, {:badnode, node}}

      nil ->
        {:exit, :noproc}
    end
  end

  defp checkout_call(pid, callers, queue?, start, timeout) do
    lock = Process.monitor(pid)
    send(pid, {:db_connection, {self(), lock}, {:checkout, callers, start, queue?}})

    receive do
      {:"ETS-TRANSFER", holder, pool, {^lock, ref}} ->
        Process.demonitor(lock, [:flush])
        {deadline, ops} = start_deadline(timeout, pool, ref, holder, start)
        :ets.update_element(holder, :conn, [{conn(:lock) + 1, lock} | ops])

        pool_ref =
          pool_ref(pool: pool, reference: ref, deadline: deadline, holder: holder, lock: lock)

        checkout_result(holder, pool_ref)

      {^lock, reply} ->
        Process.demonitor(lock, [:flush])
        reply

      {:DOWN, ^lock, _, _, reason} ->
        {:exit, reason}
    end
  end

  defp checkout_result(holder, pool_ref) do
    try do
      :ets.lookup(holder, :conn)
    rescue
      ArgumentError ->
        # Deadline could hit and be handled pool before using connection
        msg = "connection not available because deadline reached while in queue"
        {:error, DBConnection.ConnectionError.exception(msg)}
    else
      [conn(module: mod, state: state)] ->
        {:ok, pool_ref, mod, state}
    end
  end

  defp no_holder(holder) do
    "#{inspect(__MODULE__)} #{inspect(holder)} does not exist, pool inconsistent, please report at https://github.com/elixir-ecto/db_connection/issues"
  end

  defp holder_apply(holder, module, fun, args) do
    try do
      apply(module, fun, args)
    catch
      kind, reason ->
        {:catch, kind, reason, System.stacktrace()}
    else
      result when is_tuple(result) ->
        state = :erlang.element(:erlang.tuple_size(result), result)

        # This means a disconnect happened from the connection side
        # but, since we succeed, we can return the result and the client
        # will notice the disconnect anyway.
        try do
          :ets.update_element(holder, :conn, {conn(:state) + 1, state})
          result
        rescue
          ArgumentError ->
            augment_disconnect(result)
        end

      # If it is not a tuple, we just return it as is so we raise bad return.
      result ->
        result
    end
  end

  defp augment_disconnect({:disconnect, %DBConnection.ConnectionError{} = err, state}) do
    %{message: message} = err

    message =
      message <>
        " (the connection was closed by the pool, " <>
        "possibly due to a timeout or because the pool has been terminated)"

    {:disconnect, %{err | message: message}, state}
  end

  defp augment_disconnect(result), do: result

  defp done(pool_ref, ops, tag, info) do
    pool_ref(pool: pool, reference: ref, deadline: deadline, holder: holder) = pool_ref
    cancel_deadline(deadline)

    try do
      :ets.update_element(holder, :conn, [{conn(:deadline) + 1, nil} | ops])
      :ets.give_away(holder, pool, {tag, ref, info})
    rescue
      ArgumentError -> :ok
    else
      true -> :ok
    end
  end

  defp handle_done(holder, stop, err) do
    :ets.lookup(holder, :conn)
  rescue
    ArgumentError ->
      false
  else
    [conn(connection: pid, deadline: deadline, state: state)] ->
      cancel_deadline(deadline)
      :ets.delete(holder)
      stop.({pid, holder}, err, state)
      true
  end

  defp abs_timeout(now, opts) do
    case Keyword.get(opts, :timeout, @timeout) do
      :infinity -> Keyword.get(opts, :deadline)
      timeout -> min(now + timeout, Keyword.get(opts, :deadline))
    end
  end

  defp start_deadline(nil, _, _, _, _) do
    {nil, []}
  end

  defp start_deadline(timeout, pid, ref, holder, start) do
    deadline =
      :erlang.start_timer(timeout, pid, {ref, holder, self(), timeout - start}, abs: true)

    {deadline, [{conn(:deadline) + 1, deadline}]}
  end

  defp cancel_deadline(nil) do
    :ok
  end

  defp cancel_deadline(deadline) do
    :erlang.cancel_timer(deadline, async: true, info: false)
  end
end