git.squeep.com Git - akkoma/blob - lib/pleroma/gun/connection_pool.ex

   1 defmodule Pleroma.Gun.ConnectionPool do
   2   @registry __MODULE__
   3
   4   alias Pleroma.Gun.ConnectionPool.WorkerSupervisor
   5
   6   def children do
   7     [
   8       {Registry, keys: :unique, name: @registry},
   9       Pleroma.Gun.ConnectionPool.WorkerSupervisor
  10     ]
  11   end
  12
  13   def get_conn(uri, opts) do
  14     key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
  15
  16     case Registry.lookup(@registry, key) do
  17       # The key has already been registered, but connection is not up yet
  18       [{worker_pid, {nil, _used_by, _crf, _last_reference}}] ->
  19         get_gun_pid_from_worker(worker_pid)
  20
  21       [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
  22         GenServer.cast(worker_pid, {:add_client, self(), false})
  23         {:ok, gun_pid}
  24
  25       [] ->
  26         # :gun.set_owner fails in :connected state for whatevever reason,
  27         # so we open the connection in the process directly and send it's pid back
  28         # We trust gun to handle timeouts by itself
  29         case WorkerSupervisor.start_worker([uri, key, opts, self()]) do
  30           {:ok, _worker_pid} ->
  31             receive do
  32               {:conn_pid, pid} -> {:ok, pid}
  33             end
  34
  35           {:error, {:error, {:already_registered, worker_pid}}} ->
  36             get_gun_pid_from_worker(worker_pid)
  37
  38           err ->
  39             err
  40         end
  41     end
  42   end
  43
  44   defp get_gun_pid_from_worker(worker_pid) do
  45     # GenServer.call will block the process for timeout length if
  46     # the server crashes on startup (which will happen if gun fails to connect)
  47     # so instead we use cast + monitor
  48
  49     ref = Process.monitor(worker_pid)
  50     GenServer.cast(worker_pid, {:add_client, self(), true})
  51
  52     receive do
  53       {:conn_pid, pid} -> {:ok, pid}
  54       {:DOWN, ^ref, :process, ^worker_pid, reason} -> reason
  55     end
  56   end
  57
  58   def release_conn(conn_pid) do
  59     query_result =
  60       Registry.select(@registry, [
  61         {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
  62       ])
  63
  64     case query_result do
  65       [worker_pid] ->
  66         GenServer.cast(worker_pid, {:remove_client, self()})
  67
  68       [] ->
  69         :ok
  70     end
  71   end
  72 end