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]) do
  13     with {:ok, conn_pid} <- Gun.Conn.open(uri, opts),
  14          Process.link(conn_pid) do
  15       time = :os.system_time(:second)
  16
  17       {_, _} =
  18         Registry.update_value(@registry, key, fn _ ->
  19           {conn_pid, [client_pid], 1, time}
  20         end)
  21
  22       send(client_pid, {:conn_pid, conn_pid})
  23       {:ok, %{key: key, timer: nil}, :hibernate}
  24     else
  25       err -> {:stop, err}
  26     end
  27   end
  28
  29   @impl true
  30   def handle_cast({:add_client, client_pid, send_pid_back}, %{key: key} = state) do
  31     time = :os.system_time(:second)
  32
  33     {{conn_pid, _, _, _}, _} =
  34       Registry.update_value(@registry, key, fn {conn_pid, used_by, crf, last_reference} ->
  35         {conn_pid, [client_pid | used_by], crf(time - last_reference, crf), time}
  36       end)
  37
  38     if send_pid_back, do: send(client_pid, {:conn_pid, conn_pid})
  39
  40     state =
  41       if state.timer != nil do
  42         Process.cancel_timer(state[:timer])
  43         %{state | timer: nil}
  44       else
  45         state
  46       end
  47
  48     {:noreply, state, :hibernate}
  49   end
  50
  51   @impl true
  52   def handle_cast({:remove_client, client_pid}, %{key: key} = state) do
  53     {{_conn_pid, used_by, _crf, _last_reference}, _} =
  54       Registry.update_value(@registry, key, fn {conn_pid, used_by, crf, last_reference} ->
  55         {conn_pid, List.delete(used_by, client_pid), crf, last_reference}
  56       end)
  57
  58     timer =
  59       if used_by == [] do
  60         max_idle = Pleroma.Config.get([:connections_pool, :max_idle_time], 30_000)
  61         Process.send_after(self(), :idle_close, max_idle)
  62       else
  63         nil
  64       end
  65
  66     {:noreply, %{state | timer: timer}, :hibernate}
  67   end
  68
  69   @impl true
  70   def handle_info(:idle_close, state) do
  71     # Gun monitors the owner process, and will close the connection automatically
  72     # when it's terminated
  73     {:stop, :normal, state}
  74   end
  75
  76   # Gracefully shutdown if the connection got closed without any streams left
  77   @impl true
  78   def handle_info({:gun_down, _pid, _protocol, _reason, []}, state) do
  79     {:stop, :normal, state}
  80   end
  81
  82   # Otherwise, shutdown with an error
  83   @impl true
  84   def handle_info({:gun_down, _pid, _protocol, _reason, _killed_streams} = down_message, state) do
  85     {:stop, {:error, down_message}, state}
  86   end
  87
  88   # LRFU policy: https://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.55.1478
  89   defp crf(time_delta, prev_crf) do
  90     1 + :math.pow(0.5, time_delta / 100) * prev_crf
  91   end
  92 end