|
| 1 | +defmodule Plausible.Session.Persistence.TinySock do |
| 2 | + @moduledoc ~S""" |
| 3 | + Communication over Unix domain sockets. |
| 4 | +
|
| 5 | + ## Usage |
| 6 | +
|
| 7 | + ```elixir |
| 8 | + TinySock.server( |
| 9 | + base_path: "/tmp", |
| 10 | + handler: fn |
| 11 | + {"DUMP-ETS", requested_version, path} -> |
| 12 | + if requested_version == SessionV2.module_info[:md5] do |
| 13 | + for tab <- [:sessions1, :sessions2, :sessions3] do |
| 14 | + :ok = :ets.tab2file(tab, Path.join(path, "ets#{tab}")) |
| 15 | + end |
| 16 | +
|
| 17 | + :ok |
| 18 | + else |
| 19 | + {:error, :invalid_version} |
| 20 | + end |
| 21 | + end |
| 22 | + ) |
| 23 | +
|
| 24 | + dump_path = "/tmp/ysSEjw" |
| 25 | + File.mkdir_p!(dump_path) |
| 26 | + [sock_path] = TinySock.list("/tmp") |
| 27 | +
|
| 28 | + with :ok <- TinySock.call(sock_path, {"DUMP-ETS", SessionV2.module_info[:md5], dump_path}) do |
| 29 | + for "ets" <> tab <- File.ls!(dump_path) do |
| 30 | + :ets.file2tab(Path.join(dump_path, tab)) |
| 31 | + end |
| 32 | + end |
| 33 | + ``` |
| 34 | + """ |
| 35 | + |
| 36 | + use GenServer, restart: :transient |
| 37 | + require Logger |
| 38 | + |
| 39 | + @listen_opts [:binary, packet: :raw, nodelay: true, backlog: 128, active: false] |
| 40 | + @connect_opts [:binary, packet: :raw, nodelay: true, active: false] |
| 41 | + |
| 42 | + @tag_data "tinysock" |
| 43 | + @tag_data_size byte_size(@tag_data) |
| 44 | + |
| 45 | + def server(opts), do: start_link(opts) |
| 46 | + def socket(server), do: GenServer.call(server, :socket) |
| 47 | + |
| 48 | + def acceptors(server) do |
| 49 | + :ets.tab2list(GenServer.call(server, :acceptors)) |
| 50 | + end |
| 51 | + |
| 52 | + def stop(server), do: GenServer.stop(server) |
| 53 | + |
| 54 | + @doc "TODO" |
| 55 | + def list(base_path) do |
| 56 | + with {:ok, names} <- File.ls(base_path) do |
| 57 | + sock_paths = |
| 58 | + for @tag_data <> _rand = name <- names do |
| 59 | + Path.join(base_path, name) |
| 60 | + end |
| 61 | + |
| 62 | + {:ok, sock_paths} |
| 63 | + end |
| 64 | + end |
| 65 | + |
| 66 | + @doc "TODO" |
| 67 | + def call(sock_path, message, timeout \\ :timer.seconds(5)) do |
| 68 | + with {:ok, socket} <- sock_connect_or_rm(sock_path, timeout) do |
| 69 | + try do |
| 70 | + with :ok <- sock_send(socket, :erlang.term_to_binary(message)) do |
| 71 | + sock_recv(socket, timeout) |
| 72 | + end |
| 73 | + after |
| 74 | + sock_shut_and_close(socket) |
| 75 | + end |
| 76 | + end |
| 77 | + end |
| 78 | + |
| 79 | + @doc false |
| 80 | + def start_link(opts) do |
| 81 | + {gen_opts, opts} = Keyword.split(opts, [:debug, :name, :spawn_opt, :hibernate_after]) |
| 82 | + base_path = Keyword.fetch!(opts, :base_path) |
| 83 | + handler = Keyword.fetch!(opts, :handler) |
| 84 | + |
| 85 | + case File.mkdir_p(base_path) do |
| 86 | + :ok -> |
| 87 | + GenServer.start_link(__MODULE__, {base_path, handler}, gen_opts) |
| 88 | + |
| 89 | + {:error, reason} -> |
| 90 | + Logger.warning( |
| 91 | + "tinysock failed to create directory at #{inspect(base_path)}, reason: #{inspect(reason)}" |
| 92 | + ) |
| 93 | + |
| 94 | + :ignore |
| 95 | + end |
| 96 | + end |
| 97 | + |
| 98 | + @impl true |
| 99 | + def init({base_path, handler}) do |
| 100 | + case sock_listen_or_retry(base_path) do |
| 101 | + {:ok, socket} -> |
| 102 | + acceptors = :ets.new(:acceptors, [:protected]) |
| 103 | + state = {socket, acceptors, handler} |
| 104 | + for _ <- 1..10, do: spawn_acceptor(state) |
| 105 | + {:ok, state} |
| 106 | + |
| 107 | + {:error, reason} -> |
| 108 | + Logger.warning( |
| 109 | + "tinysock failed to open a listen socket in #{inspect(base_path)}, reason: #{inspect(reason)}" |
| 110 | + ) |
| 111 | + |
| 112 | + :ignore |
| 113 | + end |
| 114 | + end |
| 115 | + |
| 116 | + @impl true |
| 117 | + def handle_call(:acceptors, _from, {_socket, acceptors, _handler} = state) do |
| 118 | + {:reply, acceptors, state} |
| 119 | + end |
| 120 | + |
| 121 | + def handle_call(:socket, _from, {socket, _acceptors, _handler} = state) do |
| 122 | + {:reply, socket, state} |
| 123 | + end |
| 124 | + |
| 125 | + @impl true |
| 126 | + def handle_cast(:accepted, {socket, _acceptors, _handler} = state) do |
| 127 | + if socket, do: spawn_acceptor(state) |
| 128 | + {:noreply, state} |
| 129 | + end |
| 130 | + |
| 131 | + @impl true |
| 132 | + def handle_info({:DOWN, _ref, :process, pid, reason}, state) do |
| 133 | + case reason do |
| 134 | + :normal -> |
| 135 | + remove_acceptor(state, pid) |
| 136 | + {:noreply, state} |
| 137 | + |
| 138 | + :emfile -> |
| 139 | + raise File.Error, reason: reason, action: "accept socket", path: "tinysock lol" |
| 140 | + |
| 141 | + reason -> |
| 142 | + # :telemetry.execute([:reuse, :acceptor, :crash], reason) |
| 143 | + Logger.error("tinysock acceptor crashed, reason: #{inspect(reason)}") |
| 144 | + {:noreply, state} |
| 145 | + end |
| 146 | + end |
| 147 | + |
| 148 | + defp remove_acceptor({_socket, acceptors, _handler}, pid) do |
| 149 | + :ets.delete(acceptors, pid) |
| 150 | + end |
| 151 | + |
| 152 | + defp spawn_acceptor({socket, acceptors, handler}) do |
| 153 | + {pid, _ref} = |
| 154 | + :proc_lib.spawn_opt( |
| 155 | + __MODULE__, |
| 156 | + :accept_loop, |
| 157 | + [_parent = self(), socket, handler], |
| 158 | + [:monitor] |
| 159 | + ) |
| 160 | + |
| 161 | + :ets.insert(acceptors, {pid}) |
| 162 | + end |
| 163 | + |
| 164 | + @doc false |
| 165 | + def accept_loop(parent, listen_socket, handler) do |
| 166 | + case :gen_tcp.accept(listen_socket, :timer.seconds(5)) do |
| 167 | + {:ok, socket} -> |
| 168 | + GenServer.cast(parent, :accepted) |
| 169 | + handle_message(socket, handler) |
| 170 | + |
| 171 | + {:error, :timeout} -> |
| 172 | + accept_loop(parent, listen_socket, handler) |
| 173 | + |
| 174 | + {:error, :closed} -> |
| 175 | + :ok |
| 176 | + |
| 177 | + {:error, reason} -> |
| 178 | + exit(reason) |
| 179 | + end |
| 180 | + end |
| 181 | + |
| 182 | + defp handle_message(socket, handler) do |
| 183 | + {:ok, message} = sock_recv(socket, _timeout = :timer.seconds(5)) |
| 184 | + sock_send(socket, :erlang.term_to_binary(handler.(message))) |
| 185 | + after |
| 186 | + sock_shut_and_close(socket) |
| 187 | + end |
| 188 | + |
| 189 | + defp sock_listen_or_retry(base_path) do |
| 190 | + sock_name = @tag_data <> Base.url_encode64(:crypto.strong_rand_bytes(4), padding: false) |
| 191 | + sock_path = Path.join(base_path, sock_name) |
| 192 | + |
| 193 | + case :gen_tcp.listen(0, [{:ifaddr, {:local, sock_path}} | @listen_opts]) do |
| 194 | + {:ok, socket} -> {:ok, socket} |
| 195 | + {:error, :eaddrinuse} -> sock_listen_or_retry(base_path) |
| 196 | + {:error, reason} -> {:error, reason} |
| 197 | + end |
| 198 | + end |
| 199 | + |
| 200 | + defp sock_connect_or_rm(sock_path, timeout) do |
| 201 | + case :gen_tcp.connect({:local, sock_path}, 0, @connect_opts, timeout) do |
| 202 | + {:ok, socket} -> |
| 203 | + {:ok, socket} |
| 204 | + |
| 205 | + {:error, :timeout} = error -> |
| 206 | + error |
| 207 | + |
| 208 | + {:error, _reason} = error -> |
| 209 | + Logger.notice( |
| 210 | + "tinysock failed to connect to #{inspect(sock_path)}, reason: #{inspect(error)}" |
| 211 | + ) |
| 212 | + |
| 213 | + _ = File.rm(sock_path) |
| 214 | + error |
| 215 | + end |
| 216 | + end |
| 217 | + |
| 218 | + defp sock_send(socket, binary) do |
| 219 | + :gen_tcp.send(socket, <<@tag_data, byte_size(binary)::64, binary::bytes>>) |
| 220 | + end |
| 221 | + |
| 222 | + defp sock_recv(socket, timeout) do |
| 223 | + with {:ok, <<@tag_data, size::64>>} <- :gen_tcp.recv(socket, @tag_data_size + 8, timeout), |
| 224 | + {:ok, binary} <- :gen_tcp.recv(socket, size, timeout) do |
| 225 | + try do |
| 226 | + {:ok, :erlang.binary_to_term(binary, [:safe])} |
| 227 | + rescue |
| 228 | + e -> {:error, e} |
| 229 | + end |
| 230 | + end |
| 231 | + end |
| 232 | + |
| 233 | + defp sock_shut_and_close(socket) do |
| 234 | + :gen_tcp.shutdown(socket, :read_write) |
| 235 | + :gen_tcp.close(socket) |
| 236 | + end |
| 237 | +end |
0 commit comments