defmodule Pleroma.Gun.ConnectionPool do
@registry __MODULE__
- def get_conn(uri, opts) do
- case enforce_pool_limits() do
- :ok ->
- key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
-
- case Registry.lookup(@registry, key) do
- # The key has already been registered, but connection is not up yet
- [{worker_pid, {nil, _used_by, _crf, _last_reference}}] ->
- get_gun_pid_from_worker(worker_pid)
-
- [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
- GenServer.cast(worker_pid, {:add_client, self(), false})
- {:ok, gun_pid}
-
- [] ->
- # :gun.set_owner fails in :connected state for whatevever reason,
- # so we open the connection in the process directly and send it's pid back
- # We trust gun to handle timeouts by itself
- case GenServer.start(Pleroma.Gun.ConnectionPool.Worker, [uri, key, opts, self()],
- timeout: :infinity
- ) do
- {:ok, _worker_pid} ->
- receive do
- {:conn_pid, pid} -> {:ok, pid}
- end
-
- {:error, {:error, {:already_registered, worker_pid}}} ->
- get_gun_pid_from_worker(worker_pid)
-
- err ->
- err
- end
- end
-
- :error ->
- {:error, :pool_full}
- end
- end
-
- @enforcer_key "enforcer"
- defp enforce_pool_limits() do
- max_connections = Pleroma.Config.get([:connections_pool, :max_connections])
-
- if Registry.count(@registry) >= max_connections do
- case Registry.lookup(@registry, @enforcer_key) do
- [] ->
- pid =
- spawn(fn ->
- {:ok, _pid} = Registry.register(@registry, @enforcer_key, nil)
-
- reclaim_max =
- [:connections_pool, :reclaim_multiplier]
- |> Pleroma.Config.get()
- |> Kernel.*(max_connections)
- |> round
- |> max(1)
-
- unused_conns =
- Registry.select(
- @registry,
- [
- {{:_, :"$1", {:_, :"$2", :"$3", :"$4"}}, [{:==, :"$2", []}],
- [{{:"$1", :"$3", :"$4"}}]}
- ]
- )
+ alias Pleroma.Gun.ConnectionPool.WorkerSupervisor
- case unused_conns do
- [] ->
- exit(:pool_full)
-
- unused_conns ->
- unused_conns
- |> Enum.sort(fn {_pid1, crf1, last_reference1},
- {_pid2, crf2, last_reference2} ->
- crf1 <= crf2 and last_reference1 <= last_reference2
- end)
- |> Enum.take(reclaim_max)
- |> Enum.each(fn {pid, _, _} -> GenServer.call(pid, :idle_close) end)
- end
- end)
-
- wait_for_enforcer_finish(pid)
-
- [{pid, _}] ->
- wait_for_enforcer_finish(pid)
- end
- else
- :ok
- end
+ def children do
+ [
+ {Registry, keys: :unique, name: @registry},
+ Pleroma.Gun.ConnectionPool.WorkerSupervisor
+ ]
end
- defp wait_for_enforcer_finish(pid) do
- ref = Process.monitor(pid)
-
- receive do
- {:DOWN, ^ref, :process, ^pid, :pool_full} ->
- :error
-
- {:DOWN, ^ref, :process, ^pid, :normal} ->
- :ok
+ @spec get_conn(URI.t(), keyword()) :: {:ok, pid()} | {:error, term()}
+ def get_conn(uri, opts) do
+ key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
+
+ case Registry.lookup(@registry, key) do
+ # The key has already been registered, but connection is not up yet
+ [{worker_pid, nil}] ->
+ get_gun_pid_from_worker(worker_pid, true)
+
+ [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
+ GenServer.call(worker_pid, :add_client)
+ {:ok, gun_pid}
+
+ [] ->
+ # :gun.set_owner fails in :connected state for whatevever reason,
+ # so we open the connection in the process directly and send it's pid back
+ # We trust gun to handle timeouts by itself
+ case WorkerSupervisor.start_worker([key, uri, opts, self()]) do
+ {:ok, worker_pid} ->
+ get_gun_pid_from_worker(worker_pid, false)
+
+ {:error, {:already_started, worker_pid}} ->
+ get_gun_pid_from_worker(worker_pid, true)
+
+ err ->
+ err
+ end
end
end
- defp get_gun_pid_from_worker(worker_pid) do
+ defp get_gun_pid_from_worker(worker_pid, register) do
# GenServer.call will block the process for timeout length if
# the server crashes on startup (which will happen if gun fails to connect)
# so instead we use cast + monitor
ref = Process.monitor(worker_pid)
- GenServer.cast(worker_pid, {:add_client, self(), true})
+ if register, do: GenServer.cast(worker_pid, {:add_client, self()})
receive do
- {:conn_pid, pid} -> {:ok, pid}
- {:DOWN, ^ref, :process, ^worker_pid, reason} -> reason
+ {:conn_pid, pid} ->
+ Process.demonitor(ref)
+ {:ok, pid}
+
+ {:DOWN, ^ref, :process, ^worker_pid, reason} ->
+ case reason do
+ {:shutdown, {:error, _} = error} -> error
+ {:shutdown, error} -> {:error, error}
+ _ -> {:error, reason}
+ end
end
end
+ @spec release_conn(pid()) :: :ok
def release_conn(conn_pid) do
- [worker_pid] =
+ # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _, _, _}}} when gun_pid == conn_pid ->
+ # worker_pid end)
+ query_result =
Registry.select(@registry, [
{{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
])
- GenServer.cast(worker_pid, {:remove_client, self()})
+ case query_result do
+ [worker_pid] ->
+ GenServer.call(worker_pid, :remove_client)
+
+ [] ->
+ :ok
+ end
end
end