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}] ->
  19         get_gun_pid_from_worker(worker_pid, true)
  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([key, uri, opts, self()]) do
  30           {:ok, worker_pid} ->
  31             get_gun_pid_from_worker(worker_pid, false)
  32
  33           {:error, {:already_started, worker_pid}} ->
  34             get_gun_pid_from_worker(worker_pid, true)
  35
  36           err ->
  37             err
  38         end
  39     end
  40   end
  41
  42   defp get_gun_pid_from_worker(worker_pid, register) do
  43     # GenServer.call will block the process for timeout length if
  44     # the server crashes on startup (which will happen if gun fails to connect)
  45     # so instead we use cast + monitor
  46
  47     ref = Process.monitor(worker_pid)
  48     if register, do: GenServer.cast(worker_pid, {:add_client, self(), true})
  49
  50     receive do
  51       {:conn_pid, pid} ->
  52         Process.demonitor(ref)
  53         {:ok, pid}
  54
  55       {:DOWN, ^ref, :process, ^worker_pid, reason} ->
  56         case reason do
  57           {:shutdown, error} -> error
  58           _ -> {:error, reason}
  59         end
  60     end
  61   end
  62
  63   def release_conn(conn_pid) do
  64     # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _, _, _}}} when gun_pid == conn_pid ->
  65     #    worker_pid end)
  66     query_result =
  67       Registry.select(@registry, [
  68         {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
  69       ])
  70
  71     case query_result do
  72       [worker_pid] ->
  73         GenServer.cast(worker_pid, {:remove_client, self()})
  74
  75       [] ->
  76         :ok
  77     end
  78   end
  79 end