zf

holder.ex (13031B)

---

 defmodule DBConnection.Holder do
   @moduledoc false
   require Record
 
   @queue true
   @timeout 15000
   @time_unit 1000
 
   Record.defrecord(:conn, [:connection, :module, :state, :lock, :ts, deadline: nil, status: :ok])
   Record.defrecord(:pool_ref, [:pool, :reference, :deadline, :holder, :lock])
 
   @type t :: :ets.tid()
   @type checkin_time :: non_neg_integer() | nil
 
   ## 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, ts: System.monotonic_time())
     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)
 
     try do
       :ets.give_away(holder, pool, {:checkin, ref, System.monotonic_time()})
       {: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(), [pid], opts :: Keyword.t()) ::
               {:ok, pool_ref :: any, module, checkin_time, state :: any}
               | {:error, Exception.t()}
   def checkout(pool, callers, opts) do
     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, %DBConnection.ConnectionError{} = connection_error} = error ->
         :telemetry.execute(
           [:db_connection, :connection_error],
           %{count: 1},
           %{
             error: connection_error,
             opts: opts
           }
         )
 
         error
 
       {:error, _} = error ->
         error
 
       {:redirect, caller, proxy} ->
         case checkout(proxy, [caller], opts) do
           {:ok, _, _, _, _} = ok ->
             ok
 
           {:error, %DBConnection.ConnectionError{message: message} = exception} ->
             {:error,
              %{
                exception
                | message:
                    "could not checkout the connection owned by #{inspect(caller)}. " <>
                      "When using the sandbox, connections are shared, so this may imply " <>
                      "another process is using a connection. Reason: #{message}"
              }}
 
           {:error, _} = error ->
             error
         end
 
       {:exit, reason} ->
         exit({reason, {__MODULE__, :checkout, [pool, opts]}})
     end
   end
 
   @spec checkin(pool_ref :: any) :: :ok
   def checkin(pool_ref) do
     # 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 end of a
     # transaction.
     done(pool_ref, [{conn(:lock) + 1, nil}], :checkin, System.monotonic_time())
   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}, pid | :shared | :auto, GenServer.server()) :: :ok
   def reply_redirect(from, caller, redirect) do
     GenServer.reply(from, {:redirect, caller, 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, checkin_time) :: boolean
   def handle_checkout(holder, {pid, mref}, ref, checkin_time) do
     :ets.give_away(holder, pid, {mref, ref, checkin_time})
   rescue
     ArgumentError ->
       if Process.alive?(pid) or :ets.info(holder, :owner) != self() do
         raise ArgumentError, no_holder(holder, pid)
       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) :: true
   def handle_ping(holder) do
     :ets.lookup(holder, :conn)
   rescue
     ArgumentError ->
       raise ArgumentError, no_holder(holder, nil)
   else
     [conn(connection: conn, state: state)] ->
       DBConnection.Connection.ping({conn, holder}, state)
       :ets.delete(holder)
       true
   end
 
   @spec handle_disconnect(t, Exception.t()) :: boolean
   def handle_disconnect(holder, err) do
     handle_done(holder, &DBConnection.Connection.disconnect/3, err)
   end
 
   @spec handle_stop(t, term) :: boolean
   def handle_stop(holder, err) do
     handle_done(holder, &DBConnection.Connection.stop/3, err)
   end
 
   @spec maybe_disconnect(t, integer, non_neg_integer) :: boolean()
   def maybe_disconnect(holder, start, interval) do
     ts = :ets.lookup_element(holder, :conn, conn(:ts) + 1)
 
     cond do
       ts >= start ->
         false
 
       interval == 0 ->
         true
 
       true ->
         pid = :ets.lookup_element(holder, :conn, conn(:connection) + 1)
         System.monotonic_time() > :erlang.phash2(pid, interval) + start
     end
   rescue
     _ -> false
   else
     true ->
       opts = [message: "disconnect_all requested", severity: :info]
       handle_disconnect(holder, DBConnection.ConnectionError.exception(opts))
 
     false ->
       false
   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, checkin_time}} ->
         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, checkin_time)
 
       {^lock, reply} ->
         Process.demonitor(lock, [:flush])
         reply
 
       {:DOWN, ^lock, _, _, reason} ->
         {:exit, reason}
     end
   end
 
   defp checkout_result(holder, pool_ref, checkin_time) 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, checkin_time, state}
     end
   end
 
   defp no_holder(holder, maybe_pid) do
     reason =
       case :ets.info(holder, :owner) do
         :undefined -> "does not exist"
         ^maybe_pid -> "is being given to its current owner"
         owner when owner != self() -> "does not belong to the giving process"
         _ -> "could not be given away"
       end
 
     call_reason =
       if maybe_pid do
         "Error happened when attempting to transfer to #{inspect(maybe_pid)} " <>
           "(alive: #{Process.alive?(maybe_pid)})"
       else
         "Error happened when looking up connection"
       end
 
     """
     #{inspect(__MODULE__)} #{inspect(holder)} #{reason}, pool inconsistent.
     #{call_reason}.
 
     SELF: #{inspect(self())}
     ETS INFO: #{inspect(:ets.info(holder))}
 
     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, __STACKTRACE__}
     else
       result when is_tuple(result) ->
         state = :erlang.element(:erlang.tuple_size(result), result)
 
         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