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

   1 defmodule Pleroma.Gun.ConnectionPool.Worker do
   2   alias Pleroma.Gun
   3   use GenServer, restart: :temporary
   4
   5   @registry Pleroma.Gun.ConnectionPool
   6
   7   def start_link([key | _] = opts) do
   8     GenServer.start_link(__MODULE__, opts, name: {:via, Registry, {@registry, key}})
   9   end
  10
  11   @impl true
  12   def init([_key, _uri, _opts, _client_pid] = opts) do
  13     {:ok, nil, {:continue, {:connect, opts}}}
  14   end
  15
  16   @impl true
  17   def handle_continue({:connect, [key, uri, opts, client_pid]}, _) do
  18     with {:ok, conn_pid, protocol} <- Gun.Conn.open(uri, opts),
  19          Process.link(conn_pid) do
  20       time = :erlang.monotonic_time(:millisecond)
  21
  22       {_, _} =
  23         Registry.update_value(@registry, key, fn _ ->
  24           {conn_pid, [client_pid], 1, time}
  25         end)
  26
  27       send(client_pid, {:conn_pid, conn_pid})
  28
  29       {:noreply,
  30        %{
  31          key: key,
  32          timer: nil,
  33          client_monitors: %{client_pid => Process.monitor(client_pid)},
  34          protocol: protocol
  35        }, :hibernate}
  36     else
  37       err ->
  38         {:stop, {:shutdown, err}, nil}
  39     end
  40   end
  41
  42   @impl true
  43   def handle_cast({:add_client, client_pid}, state) do
  44     case handle_call(:add_client, {client_pid, nil}, state) do
  45       {:reply, conn_pid, state, :hibernate} ->
  46         send(client_pid, {:conn_pid, conn_pid})
  47         {:noreply, state, :hibernate}
  48     end
  49   end
  50
  51   @impl true
  52   def handle_cast({:remove_client, client_pid}, state) do
  53     case handle_call(:remove_client, {client_pid, nil}, state) do
  54       {:reply, _, state, :hibernate} ->
  55         {:noreply, state, :hibernate}
  56     end
  57   end
  58
  59   @impl true
  60   def handle_call(:add_client, {client_pid, _}, %{key: key, protocol: protocol} = state) do
  61     time = :erlang.monotonic_time(:millisecond)
  62
  63     {{conn_pid, used_by, _, _}, _} =
  64       Registry.update_value(@registry, key, fn {conn_pid, used_by, crf, last_reference} ->
  65         {conn_pid, [client_pid | used_by], crf(time - last_reference, crf), time}
  66       end)
  67
  68     :telemetry.execute(
  69       [:pleroma, :connection_pool, :client, :add],
  70       %{client_pid: client_pid, clients: used_by},
  71       %{key: state.key, protocol: protocol}
  72     )
  73
  74     state =
  75       if state.timer != nil do
  76         Process.cancel_timer(state[:timer])
  77         %{state | timer: nil}
  78       else
  79         state
  80       end
  81
  82     ref = Process.monitor(client_pid)
  83
  84     state = put_in(state.client_monitors[client_pid], ref)
  85     {:reply, conn_pid, state, :hibernate}
  86   end
  87
  88   @impl true
  89   def handle_call(:remove_client, {client_pid, _}, %{key: key} = state) do
  90     {{_conn_pid, used_by, _crf, _last_reference}, _} =
  91       Registry.update_value(@registry, key, fn {conn_pid, used_by, crf, last_reference} ->
  92         {conn_pid, List.delete(used_by, client_pid), crf, last_reference}
  93       end)
  94
  95     {ref, state} = pop_in(state.client_monitors[client_pid])
  96     # DOWN message can receive right after `remove_client` call and cause worker to terminate
  97     state =
  98       if is_nil(ref) do
  99         state
 100       else
 101         Process.demonitor(ref)
 102
 103         timer =
 104           if used_by == [] do
 105             max_idle = Pleroma.Config.get([:connections_pool, :max_idle_time], 30_000)
 106             Process.send_after(self(), :idle_close, max_idle)
 107           else
 108             nil
 109           end
 110
 111         %{state | timer: timer}
 112       end
 113
 114     {:reply, :ok, state, :hibernate}
 115   end
 116
 117   @impl true
 118   def handle_info(:idle_close, state) do
 119     # Gun monitors the owner process, and will close the connection automatically
 120     # when it's terminated
 121     {:stop, :normal, state}
 122   end
 123
 124   @impl true
 125   def handle_info({:gun_up, _pid, _protocol}, state) do
 126     {:noreply, state, :hibernate}
 127   end
 128
 129   # Gracefully shutdown if the connection got closed without any streams left
 130   @impl true
 131   def handle_info({:gun_down, _pid, _protocol, _reason, []}, state) do
 132     {:stop, :normal, state}
 133   end
 134
 135   # Otherwise, wait for retry
 136   @impl true
 137   def handle_info({:gun_down, _pid, _protocol, _reason, _killed_streams}, state) do
 138     {:noreply, state, :hibernate}
 139   end
 140
 141   @impl true
 142   def handle_info({:DOWN, _ref, :process, pid, reason}, state) do
 143     :telemetry.execute(
 144       [:pleroma, :connection_pool, :client, :dead],
 145       %{client_pid: pid, reason: reason},
 146       %{key: state.key}
 147     )
 148
 149     handle_cast({:remove_client, pid}, state)
 150   end
 151
 152   # LRFU policy: https://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.55.1478
 153   defp crf(time_delta, prev_crf) do
 154     1 + :math.pow(0.5, 0.0001 * time_delta) * prev_crf
 155   end
 156 end