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

   1 defmodule Pleroma.Gun.ConnectionPool.WorkerSupervisor do
   2   @moduledoc "Supervisor for pool workers. Does not do anything except enforce max connection limit"
   3
   4   use DynamicSupervisor
   5
   6   def start_link(opts) do
   7     DynamicSupervisor.start_link(__MODULE__, opts, name: __MODULE__)
   8   end
   9
  10   def init(_opts) do
  11     DynamicSupervisor.init(
  12       strategy: :one_for_one,
  13       max_children: Pleroma.Config.get([:connections_pool, :max_connections])
  14     )
  15   end
  16
  17   def start_worker(opts) do
  18     case DynamicSupervisor.start_child(__MODULE__, {Pleroma.Gun.ConnectionPool.Worker, opts}) do
  19       {:error, :max_children} ->
  20         case free_pool() do
  21           :ok -> start_worker(opts)
  22           :error -> {:error, :pool_full}
  23         end
  24
  25       res ->
  26         res
  27     end
  28   end
  29
  30   @registry Pleroma.Gun.ConnectionPool
  31   @enforcer_key "enforcer"
  32   defp free_pool do
  33     case Registry.lookup(@registry, @enforcer_key) do
  34       [] ->
  35         pid =
  36           spawn(fn ->
  37             {:ok, _pid} = Registry.register(@registry, @enforcer_key, nil)
  38             max_connections = Pleroma.Config.get([:connections_pool, :max_connections])
  39
  40             reclaim_max =
  41               [:connections_pool, :reclaim_multiplier]
  42               |> Pleroma.Config.get()
  43               |> Kernel.*(max_connections)
  44               |> round
  45               |> max(1)
  46
  47             unused_conns =
  48               Registry.select(
  49                 @registry,
  50                 [
  51                   {{:_, :"$1", {:_, :"$2", :"$3", :"$4"}}, [{:==, :"$2", []}],
  52                    [{{:"$1", :"$3", :"$4"}}]}
  53                 ]
  54               )
  55
  56             case unused_conns do
  57               [] ->
  58                 exit(:no_unused_conns)
  59
  60               unused_conns ->
  61                 unused_conns
  62                 |> Enum.sort(fn {_pid1, crf1, last_reference1}, {_pid2, crf2, last_reference2} ->
  63                   crf1 <= crf2 and last_reference1 <= last_reference2
  64                 end)
  65                 |> Enum.take(reclaim_max)
  66                 |> Enum.each(fn {pid, _, _} ->
  67                   DynamicSupervisor.terminate_child(__MODULE__, pid)
  68                 end)
  69             end
  70           end)
  71
  72         wait_for_enforcer_finish(pid)
  73
  74       [{pid, _}] ->
  75         wait_for_enforcer_finish(pid)
  76     end
  77   end
  78
  79   defp wait_for_enforcer_finish(pid) do
  80     ref = Process.monitor(pid)
  81
  82     receive do
  83       {:DOWN, ^ref, :process, ^pid, :no_unused_conns} ->
  84         :error
  85
  86       {:DOWN, ^ref, :process, ^pid, :normal} ->
  87         :ok
  88     end
  89   end
  90 end