zf

manager.ex (10437B)

---

 defmodule DBConnection.Ownership.Manager do
   @moduledoc false
   use GenServer
   require Logger
   alias DBConnection.Ownership.Proxy
 
   @timeout 5_000
 
   @callback start_link({module, opts :: Keyword.t()}) ::
               GenServer.on_start()
   def start_link({module, opts}) do
     {owner_opts, pool_opts} = Keyword.split(opts, [:name])
     GenServer.start_link(__MODULE__, {module, owner_opts, pool_opts}, owner_opts)
   end
 
   @callback disconnect_all(GenServer.server(), non_neg_integer, Keyword.t()) :: :ok
   def disconnect_all(pool, interval, opts) do
     inner_pool = GenServer.call(pool, :pool, :infinity)
     DBConnection.ConnectionPool.disconnect_all(inner_pool, interval, opts)
   end
 
   @spec proxy_for(callers :: [pid], Keyword.t()) :: {caller :: pid, proxy :: pid} | nil
   def proxy_for(callers, opts) do
     case Keyword.fetch(opts, :name) do
       {:ok, name} ->
         Enum.find_value(callers, &List.first(:ets.lookup(name, &1)))
 
       :error ->
         nil
     end
   end
 
   @spec checkout(GenServer.server(), Keyword.t()) ::
           {:ok, pid} | {:already, :owner | :allowed}
   def checkout(manager, opts) do
     GenServer.call(manager, {:checkout, opts}, :infinity)
   end
 
   @spec checkin(GenServer.server(), Keyword.t()) ::
           :ok | :not_owner | :not_found
   def checkin(manager, opts) do
     timeout = Keyword.get(opts, :timeout, @timeout)
     GenServer.call(manager, :checkin, timeout)
   end
 
   @spec mode(GenServer.server(), :auto | :manual | {:shared, pid}, Keyword.t()) ::
           :ok | :already_shared | :not_owner | :not_found
   def mode(manager, mode, opts)
       when mode in [:auto, :manual]
       when elem(mode, 0) == :shared and is_pid(elem(mode, 1)) do
     timeout = Keyword.get(opts, :timeout, @timeout)
     GenServer.call(manager, {:mode, mode}, timeout)
   end
 
   @spec allow(GenServer.server(), parent :: pid, allow :: pid, Keyword.t()) ::
           :ok | {:already, :owner | :allowed} | :not_found
   def allow(manager, parent, allow, opts) do
     timeout = Keyword.get(opts, :timeout, @timeout)
     GenServer.call(manager, {:allow, parent, allow}, timeout)
   end
 
   ## Callbacks
 
   def init({module, owner_opts, pool_opts}) do
     DBConnection.register_as_pool(module)
 
     ets =
       case Keyword.fetch(owner_opts, :name) do
         {:ok, name} when is_atom(name) ->
           :ets.new(name, [:set, :named_table, :protected, read_concurrency: true])
 
         _ ->
           nil
       end
 
     # We can only start the connection pool directly because
     # neither the pool's GenServer nor the manager trap exits.
     # Otherwise we would need a supervisor plus a watcher process.
     pool_opts = Keyword.delete(pool_opts, :pool)
     {:ok, pool} = DBConnection.start_link(module, pool_opts)
 
     log = Keyword.get(pool_opts, :ownership_log, nil)
     mode = Keyword.get(pool_opts, :ownership_mode, :auto)
     checkout_opts = Keyword.take(pool_opts, [:ownership_timeout, :queue_target, :queue_interval])
 
     {:ok,
      %{
        pool: pool,
        checkouts: %{},
        owners: %{},
        checkout_opts: checkout_opts,
        mode: mode,
        mode_ref: nil,
        ets: ets,
        log: log
      }}
   end
 
   def handle_call(:pool, _from, %{pool: pool} = state) do
     {:reply, pool, state}
   end
 
   def handle_call({:mode, {:shared, shared}}, {caller, _}, %{mode: {:shared, current}} = state) do
     cond do
       shared == current ->
         {:reply, :ok, state}
 
       Process.alive?(current) ->
         {:reply, :already_shared, state}
 
       true ->
         share_and_reply(state, shared, caller)
     end
   end
 
   def handle_call({:mode, {:shared, shared}}, {caller, _}, state) do
     share_and_reply(state, shared, caller)
   end
 
   def handle_call({:mode, mode}, _from, %{mode: mode} = state) do
     {:reply, :ok, state}
   end
 
   def handle_call({:mode, mode}, {caller, _}, state) do
     state = proxy_checkin_all_except(state, [], caller)
     {:reply, :ok, %{state | mode: mode, mode_ref: nil}}
   end
 
   def handle_call(:checkin, {caller, _}, state) do
     {reply, state} = proxy_checkin(state, caller, caller)
     {:reply, reply, state}
   end
 
   def handle_call({:allow, caller, allow}, _from, %{checkouts: checkouts} = state) do
     if kind = already_checked_out(checkouts, allow) do
       {:reply, {:already, kind}, state}
     else
       case Map.get(checkouts, caller, :not_found) do
         {:owner, ref, proxy} ->
           {:reply, :ok, owner_allow(state, allow, ref, proxy)}
 
         {:allowed, ref, proxy} ->
           {:reply, :ok, owner_allow(state, allow, ref, proxy)}
 
         :not_found ->
           {:reply, :not_found, state}
       end
     end
   end
 
   def handle_call({:checkout, opts}, {caller, _}, %{checkouts: checkouts} = state) do
     if kind = already_checked_out(checkouts, caller) do
       {:reply, {:already, kind}, state}
     else
       {proxy, state} = proxy_checkout(state, caller, opts)
       {:reply, {:ok, proxy}, state}
     end
   end
 
   def handle_info({:db_connection, from, {:checkout, callers, _now, queue?}}, state) do
     %{checkouts: checkouts, mode: mode, checkout_opts: checkout_opts} = state
     caller = find_caller(callers, checkouts, mode)
 
     case Map.get(checkouts, caller, :not_found) do
       {status, _ref, proxy} when status in [:owner, :allowed] ->
         DBConnection.Holder.reply_redirect(from, caller, proxy)
         {:noreply, state}
 
       :not_found when mode == :auto ->
         {proxy, state} = proxy_checkout(state, caller, [queue: queue?] ++ checkout_opts)
         DBConnection.Holder.reply_redirect(from, caller, proxy)
         {:noreply, state}
 
       :not_found when mode == :manual ->
         not_found(from)
         {:noreply, state}
 
       :not_found ->
         {:shared, shared} = mode
         {:owner, _ref, proxy} = Map.fetch!(checkouts, shared)
         DBConnection.Holder.reply_redirect(from, shared, proxy)
         {:noreply, state}
     end
   end
 
   def handle_info({:DOWN, ref, _, _, _}, state) do
     {:noreply, state |> owner_down(ref) |> unshare(ref)}
   end
 
   def handle_info(_msg, state) do
     {:noreply, state}
   end
 
   defp already_checked_out(checkouts, pid) do
     case Map.get(checkouts, pid, :not_found) do
       {:owner, _, _} -> :owner
       {:allowed, _, _} -> :allowed
       :not_found -> nil
     end
   end
 
   defp proxy_checkout(state, caller, opts) do
     %{pool: pool, checkouts: checkouts, owners: owners, ets: ets, log: log} = state
 
     {:ok, proxy} =
       DynamicSupervisor.start_child(
         DBConnection.Ownership.Supervisor,
         {DBConnection.Ownership.Proxy, {caller, pool, opts}}
       )
 
     log && Logger.log(log, fn -> [inspect(caller), " owns proxy " | inspect(proxy)] end)
     ref = Process.monitor(proxy)
     checkouts = Map.put(checkouts, caller, {:owner, ref, proxy})
     owners = Map.put(owners, ref, {proxy, caller, []})
     ets && :ets.insert(ets, {caller, proxy})
     {proxy, %{state | checkouts: checkouts, owners: owners}}
   end
 
   defp proxy_checkin(state, maybe_owner, caller) do
     case get_and_update_in(state.checkouts, &Map.pop(&1, maybe_owner, :not_found)) do
       {{:owner, ref, proxy}, state} ->
         Proxy.stop(proxy, caller)
         {:ok, state |> owner_down(ref) |> unshare(ref)}
 
       {{:allowed, _, _}, _} ->
         {:not_owner, state}
 
       {:not_found, _} ->
         {:not_found, state}
     end
   end
 
   defp proxy_checkin_all_except(state, except, caller) do
     Enum.reduce(state.checkouts, state, fn {pid, _}, state ->
       if pid in except do
         state
       else
         {_, state} = proxy_checkin(state, pid, caller)
         state
       end
     end)
   end
 
   defp owner_allow(%{ets: ets, log: log} = state, allow, ref, proxy) do
     log && Logger.log(log, fn -> [inspect(allow), " allowed on proxy " | inspect(proxy)] end)
     state = put_in(state.checkouts[allow], {:allowed, ref, proxy})
 
     state =
       update_in(state.owners[ref], fn {proxy, caller, allowed} ->
         {proxy, caller, [allow | List.delete(allowed, allow)]}
       end)
 
     ets && :ets.insert(ets, {allow, proxy})
     state
   end
 
   defp owner_down(%{ets: ets, log: log} = state, ref) do
     case get_and_update_in(state.owners, &Map.pop(&1, ref)) do
       {{proxy, caller, allowed}, state} ->
         Process.demonitor(ref, [:flush])
         entries = [caller | allowed]
 
         log &&
           Logger.log(log, fn ->
             [Enum.map_join(entries, ", ", &inspect/1), " lose proxy " | inspect(proxy)]
           end)
 
         ets && Enum.each(entries, &:ets.delete(ets, &1))
         update_in(state.checkouts, &Map.drop(&1, entries))
 
       {nil, state} ->
         state
     end
   end
 
   defp share_and_reply(%{checkouts: checkouts} = state, shared, caller) do
     case Map.get(checkouts, shared, :not_found) do
       {:owner, ref, _} ->
         state = proxy_checkin_all_except(state, [shared], caller)
         {:reply, :ok, %{state | mode: {:shared, shared}, mode_ref: ref}}
 
       {:allowed, _, _} ->
         {:reply, :not_owner, state}
 
       :not_found ->
         {:reply, :not_found, state}
     end
   end
 
   defp unshare(%{mode_ref: ref} = state, ref) do
     %{state | mode: :manual, mode_ref: nil}
   end
 
   defp unshare(state, _ref) do
     state
   end
 
   defp find_caller(callers, checkouts, :manual) do
     Enum.find(callers, &Map.has_key?(checkouts, &1)) || hd(callers)
   end
 
   defp find_caller([caller | _], _checkouts, _mode) do
     caller
   end
 
   defp not_found({pid, _} = from) do
     msg = """
     cannot find ownership process for #{inspect(pid)}.
 
     When using ownership, you must manage connections in one
     of the four ways:
 
     * By explicitly checking out a connection
     * By explicitly allowing a spawned process
     * By running the pool in shared mode
     * By using :caller option with allowed process
 
     The first two options require every new process to explicitly
     check a connection out or be allowed by calling checkout or
     allow respectively.
 
     The third option requires a {:shared, pid} mode to be set.
     If using shared mode in tests, make sure your tests are not
     async.
 
     The fourth option requires [caller: pid] to be used when
     checking out a connection from the pool. The caller process
     should already be allowed on a connection.
 
     If you are reading this error, it means you have not done one
     of the steps above or that the owner process has crashed.
     """
 
     DBConnection.Holder.reply_error(from, DBConnection.OwnershipError.exception(msg))
   end
 end