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   @spec get_conn(URI.t(), keyword()) :: {:ok, pid()} | {:error, term()}
  14   def get_conn(uri, opts) do
  15     key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
  16
  17     case Registry.lookup(@registry, key) do
  18       # The key has already been registered, but connection is not up yet
  19       [{worker_pid, nil}] ->
  20         get_gun_pid_from_worker(worker_pid, true)
  21
  22       [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
  23         GenServer.call(worker_pid, :add_client)
  24         {:ok, gun_pid}
  25
  26       [] ->
  27         # :gun.set_owner fails in :connected state for whatevever reason,
  28         # so we open the connection in the process directly and send it's pid back
  29         # We trust gun to handle timeouts by itself
  30         case WorkerSupervisor.start_worker([key, uri, opts, self()]) do
  31           {:ok, worker_pid} ->
  32             get_gun_pid_from_worker(worker_pid, false)
  33
  34           {:error, {:already_started, worker_pid}} ->
  35             get_gun_pid_from_worker(worker_pid, true)
  36
  37           err ->
  38             err
  39         end
  40     end
  41   end
  42
  43   defp get_gun_pid_from_worker(worker_pid, register) do
  44     # GenServer.call will block the process for timeout length if
  45     # the server crashes on startup (which will happen if gun fails to connect)
  46     # so instead we use cast + monitor
  47
  48     ref = Process.monitor(worker_pid)
  49     if register, do: GenServer.cast(worker_pid, {:add_client, self()})
  50
  51     receive do
  52       {:conn_pid, pid} ->
  53         Process.demonitor(ref)
  54         {:ok, pid}
  55
  56       {:DOWN, ^ref, :process, ^worker_pid, reason} ->
  57         case reason do
  58           {:shutdown, {:error, _} = error} -> error
  59           {:shutdown, error} -> {:error, error}
  60           _ -> {:error, reason}
  61         end
  62     end
  63   end
  64
  65   @spec release_conn(pid()) :: :ok
  66   def release_conn(conn_pid) do
  67     # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _, _, _}}} when gun_pid == conn_pid ->
  68     #    worker_pid end)
  69     query_result =
  70       Registry.select(@registry, [
  71         {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
  72       ])
  73
  74     case query_result do
  75       [worker_pid] ->
  76         GenServer.call(worker_pid, :remove_client)
  77
  78       [] ->
  79         :ok
  80     end
  81   end
  82 end