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

   1 defmodule Pleroma.Gun.ConnectionPool.WorkerSupervisor do
   2   @doc "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
  39             max_connections = Pleroma.Config.get([:connections_pool, :max_connections])
  40
  41             reclaim_max =
  42               [:connections_pool, :reclaim_multiplier]
  43               |> Pleroma.Config.get()
  44               |> Kernel.*(max_connections)
  45               |> round
  46               |> max(1)
  47
  48             unused_conns =
  49               Registry.select(
  50                 @registry,
  51                 [
  52                   {{:_, :"$1", {:_, :"$2", :"$3", :"$4"}}, [{:==, :"$2", []}],
  53                    [{{:"$1", :"$3", :"$4"}}]}
  54                 ]
  55               )
  56
  57             case unused_conns do
  58               [] ->
  59                 exit(:no_unused_conns)
  60
  61               unused_conns ->
  62                 unused_conns
  63                 |> Enum.sort(fn {_pid1, crf1, last_reference1}, {_pid2, crf2, last_reference2} ->
  64                   crf1 <= crf2 and last_reference1 <= last_reference2
  65                 end)
  66                 |> Enum.take(reclaim_max)
  67                 |> Enum.each(fn {pid, _, _} ->
  68                   DynamicSupervisor.terminate_child(__MODULE__, pid)
  69                 end)
  70             end
  71           end)
  72
  73         wait_for_enforcer_finish(pid)
  74
  75       [{pid, _}] ->
  76         wait_for_enforcer_finish(pid)
  77     end
  78   end
  79
  80   defp wait_for_enforcer_finish(pid) do
  81     ref = Process.monitor(pid)
  82
  83     receive do
  84       {:DOWN, ^ref, :process, ^pid, :no_unused_conns} ->
  85         :error
  86
  87       {:DOWN, ^ref, :process, ^pid, :normal} ->
  88         :ok
  89     end
  90   end
  91 end