Connection pool: Fix race conditions in limit enforcement
[akkoma] / lib / pleroma / gun / connection_pool.ex
1 defmodule Pleroma.Gun.ConnectionPool do
2 @registry __MODULE__
3
4 alias Pleroma.Gun.ConnectionPool.WorkerSupervisor
5
6 def children do
7 [
8 {Registry, keys: :unique, name: @registry},
9 Pleroma.Gun.ConnectionPool.WorkerSupervisor
10 ]
11 end
12
13 def get_conn(uri, opts) do
14 key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
15
16 case Registry.lookup(@registry, key) do
17 # The key has already been registered, but connection is not up yet
18 [{worker_pid, {nil, _used_by, _crf, _last_reference}}] ->
19 get_gun_pid_from_worker(worker_pid)
20
21 [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
22 GenServer.cast(worker_pid, {:add_client, self(), false})
23 {:ok, gun_pid}
24
25 [] ->
26 # :gun.set_owner fails in :connected state for whatevever reason,
27 # so we open the connection in the process directly and send it's pid back
28 # We trust gun to handle timeouts by itself
29 case WorkerSupervisor.start_worker([uri, key, opts, self()]) do
30 {:ok, _worker_pid} ->
31 receive do
32 {:conn_pid, pid} -> {:ok, pid}
33 end
34
35 {:error, {:error, {:already_registered, worker_pid}}} ->
36 get_gun_pid_from_worker(worker_pid)
37
38 err ->
39 err
40 end
41 end
42 end
43
44 defp get_gun_pid_from_worker(worker_pid) do
45 # GenServer.call will block the process for timeout length if
46 # the server crashes on startup (which will happen if gun fails to connect)
47 # so instead we use cast + monitor
48
49 ref = Process.monitor(worker_pid)
50 GenServer.cast(worker_pid, {:add_client, self(), true})
51
52 receive do
53 {:conn_pid, pid} -> {:ok, pid}
54 {:DOWN, ^ref, :process, ^worker_pid, reason} -> reason
55 end
56 end
57
58 def release_conn(conn_pid) do
59 query_result =
60 Registry.select(@registry, [
61 {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
62 ])
63
64 case query_result do
65 [worker_pid] ->
66 GenServer.cast(worker_pid, {:remove_client, self()})
67
68 [] ->
69 :ok
70 end
71 end
72 end