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
  97     Process.demonitor(ref, [:flush])
  98
  99     timer =
 100       if used_by == [] do
 101         max_idle = Pleroma.Config.get([:connections_pool, :max_idle_time], 30_000)
 102         Process.send_after(self(), :idle_close, max_idle)
 103       else
 104         nil
 105       end
 106
 107     {:reply, :ok, %{state | timer: timer}, :hibernate}
 108   end
 109
 110   @impl true
 111   def handle_info(:idle_close, state) do
 112     # Gun monitors the owner process, and will close the connection automatically
 113     # when it's terminated
 114     {:stop, :normal, state}
 115   end
 116
 117   @impl true
 118   def handle_info({:gun_up, _pid, _protocol}, state) do
 119     {:noreply, state, :hibernate}
 120   end
 121
 122   # Gracefully shutdown if the connection got closed without any streams left
 123   @impl true
 124   def handle_info({:gun_down, _pid, _protocol, _reason, []}, state) do
 125     {:stop, :normal, state}
 126   end
 127
 128   # Otherwise, wait for retry
 129   @impl true
 130   def handle_info({:gun_down, _pid, _protocol, _reason, _killed_streams}, state) do
 131     {:noreply, state, :hibernate}
 132   end
 133
 134   @impl true
 135   def handle_info({:DOWN, _ref, :process, pid, reason}, state) do
 136     :telemetry.execute(
 137       [:pleroma, :connection_pool, :client, :dead],
 138       %{client_pid: pid, reason: reason},
 139       %{key: state.key}
 140     )
 141
 142     handle_cast({:remove_client, pid}, state)
 143   end
 144
 145   # LRFU policy: https://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.55.1478
 146   defp crf(time_delta, prev_crf) do
 147     1 + :math.pow(0.5, 0.0001 * time_delta) * prev_crf
 148   end
 149 end