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

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