total rebase
[anni] / lib / pleroma / gun / connection_pool.ex
1 # Pleroma: A lightweight social networking server
2 # Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
3 # SPDX-License-Identifier: AGPL-3.0-only
4
5 defmodule Pleroma.Gun.ConnectionPool do
6   @registry __MODULE__
7
8   alias Pleroma.Gun.ConnectionPool.WorkerSupervisor
9
10   def children do
11     [
12       {Registry, keys: :unique, name: @registry},
13       Pleroma.Gun.ConnectionPool.WorkerSupervisor
14     ]
15   end
16
17   @spec get_conn(URI.t(), keyword()) :: {:ok, pid()} | {:error, term()}
18   def get_conn(uri, opts) do
19     key = "#{uri.scheme}:#{uri.host}:#{uri.port}"
20
21     case Registry.lookup(@registry, key) do
22       # The key has already been registered, but connection is not up yet
23       [{worker_pid, nil}] ->
24         get_gun_pid_from_worker(worker_pid, true)
25
26       [{worker_pid, {gun_pid, _used_by, _crf, _last_reference}}] ->
27         GenServer.call(worker_pid, :add_client)
28         {:ok, gun_pid}
29
30       [] ->
31         # :gun.set_owner fails in :connected state for whatevever reason,
32         # so we open the connection in the process directly and send it's pid back
33         # We trust gun to handle timeouts by itself
34         case WorkerSupervisor.start_worker([key, uri, opts, self()]) do
35           {:ok, worker_pid} ->
36             get_gun_pid_from_worker(worker_pid, false)
37
38           {:error, {:already_started, worker_pid}} ->
39             get_gun_pid_from_worker(worker_pid, true)
40
41           err ->
42             err
43         end
44     end
45   end
46
47   defp get_gun_pid_from_worker(worker_pid, register) do
48     # GenServer.call will block the process for timeout length if
49     # the server crashes on startup (which will happen if gun fails to connect)
50     # so instead we use cast + monitor
51
52     ref = Process.monitor(worker_pid)
53     if register, do: GenServer.cast(worker_pid, {:add_client, self()})
54
55     receive do
56       {:conn_pid, pid} ->
57         Process.demonitor(ref)
58         {:ok, pid}
59
60       {:DOWN, ^ref, :process, ^worker_pid, reason} ->
61         case reason do
62           {:shutdown, {:error, _} = error} -> error
63           {:shutdown, error} -> {:error, error}
64           _ -> {:error, reason}
65         end
66     end
67   end
68
69   @spec release_conn(pid()) :: :ok
70   def release_conn(conn_pid) do
71     # :ets.fun2ms(fn {_, {worker_pid, {gun_pid, _, _, _}}} when gun_pid == conn_pid ->
72     #    worker_pid end)
73     query_result =
74       Registry.select(@registry, [
75         {{:_, :"$1", {:"$2", :_, :_, :_}}, [{:==, :"$2", conn_pid}], [:"$1"]}
76       ])
77
78     case query_result do
79       [worker_pid] ->
80         GenServer.call(worker_pid, :remove_client)
81
82       [] ->
83         :ok
84     end
85   end
86 end