Elixir ATProtocol ingestion and sync library.
6
fork

Configure Feed

Select the types of activity you want to include in your feed.

refactor: split core websocket functionality out

+450 -128
+6 -1
CHANGELOG.md
··· 8 8 9 9 ## [Unreleased] 10 10 11 - ### Breaking Change 11 + ### Breaking Changes 12 12 13 13 - Existing behaviour moved to `Drinkup.Firehose` namespace, to make way for 14 14 alternate sync systems. 15 + 16 + ### Changed 17 + 18 + - Refactor core connection logic for websockets into `Drinkup.Socket` to make it 19 + easy to use across multiple different services. 15 20 16 21 ## [0.1.0] - 2025-05-26 17 22
+63 -8
lib/firehose/options.ex
··· 1 1 defmodule Drinkup.Firehose.Options do 2 + @moduledoc """ 3 + Configuration options for ATProto Firehose relay subscriptions. 4 + 5 + This module defines the configuration structure for connecting to and 6 + consuming events from an ATProto Firehose relay. The Firehose streams 7 + real-time repository events from the AT Protocol network. 8 + 9 + ## Options 10 + 11 + - `:consumer` (required) - Module implementing `Drinkup.Firehose.Consumer` behaviour 12 + - `:name` - Unique name for this Firehose instance in the supervision tree (default: `Drinkup.Firehose`) 13 + - `:host` - Firehose relay URL (default: `"https://bsky.network"`) 14 + - `:cursor` - Optional sequence number to resume streaming from 15 + 16 + ## Example 17 + 18 + %{ 19 + consumer: MyFirehoseConsumer, 20 + name: MyFirehose, 21 + host: "https://bsky.network", 22 + cursor: 12345 23 + } 24 + """ 25 + 2 26 use TypedStruct 3 27 4 28 @default_host "https://bsky.network" 5 29 30 + @typedoc """ 31 + Map of configuration options accepted by `Drinkup.Firehose.child_spec/1`. 32 + """ 6 33 @type options() :: %{ 7 - required(:consumer) => module(), 8 - optional(:name) => atom(), 9 - optional(:host) => String.t(), 10 - optional(:cursor) => pos_integer() 34 + required(:consumer) => consumer(), 35 + optional(:name) => name(), 36 + optional(:host) => host(), 37 + optional(:cursor) => cursor() 11 38 } 12 39 40 + @typedoc """ 41 + Module implementing the `Drinkup.Firehose.Consumer` behaviour. 42 + """ 43 + @type consumer() :: module() 44 + 45 + @typedoc """ 46 + Unique identifier for this Firehose instance in the supervision tree. 47 + 48 + Used for Registry lookups and naming child processes. 49 + """ 50 + @type name() :: atom() 51 + 52 + @typedoc """ 53 + HTTP/HTTPS URL of the ATProto Firehose relay. 54 + 55 + Defaults to `"https://bsky.network"` which is the public Bluesky relay. 56 + """ 57 + @type host() :: String.t() 58 + 59 + @typedoc """ 60 + Optional sequence number to resume streaming from. 61 + 62 + When provided, the Firehose will replay events starting from this sequence 63 + number. Useful for resuming after a restart without missing events. The 64 + cursor is automatically tracked and updated as events are received. 65 + """ 66 + @type cursor() :: pos_integer() | nil 67 + 13 68 typedstruct do 14 - field :consumer, module(), enforce: true 15 - field :name, atom(), default: Drinkup 16 - field :host, String.t(), default: @default_host 17 - field :cursor, pos_integer() | nil 69 + field :consumer, consumer(), enforce: true 70 + field :name, name(), default: Drinkup.Firehose 71 + field :host, host(), default: @default_host 72 + field :cursor, cursor() 18 73 end 19 74 20 75 @spec from(options()) :: t()
+40 -119
lib/firehose/socket.ex
··· 1 1 defmodule Drinkup.Firehose.Socket do 2 2 @moduledoc """ 3 - gen_statem process for managing the websocket connection to an ATProto relay. 3 + WebSocket connection handler for ATProto relay subscriptions. 4 + 5 + Implements the Drinkup.Socket behaviour to manage connections to an ATProto 6 + Firehose relay, handling CAR/CBOR-encoded frames and dispatching events to 7 + the configured consumer. 4 8 """ 9 + 10 + use Drinkup.Socket 5 11 6 12 require Logger 7 13 alias Drinkup.Firehose.{Event, Options} 8 14 9 - @behaviour :gen_statem 10 - @timeout :timer.seconds(5) 11 - # TODO: `flow` determines messages in buffer. Determine ideal value? 12 - @flow 10 13 - 14 15 @op_regular 1 15 16 @op_error -1 16 - 17 - defstruct [:options, :seq, :conn, :stream] 18 17 19 18 @impl true 20 - def callback_mode, do: [:state_functions, :state_enter] 21 - 22 - def child_spec(opts) do 23 - %{ 24 - id: __MODULE__, 25 - start: {__MODULE__, :start_link, [opts, []]}, 26 - type: :worker, 27 - restart: :permanent, 28 - shutdown: 500 29 - } 19 + def init(opts) do 20 + options = Keyword.fetch!(opts, :options) 21 + {:ok, %{seq: options.cursor, options: options, host: options.host}} 30 22 end 31 23 32 24 def start_link(%Options{} = options, statem_opts) do 33 - :gen_statem.start_link(__MODULE__, options, statem_opts) 34 - end 35 - 36 - @impl true 37 - def init(%{cursor: seq} = options) do 38 - data = %__MODULE__{seq: seq, options: options} 39 - {:ok, :disconnected, data, [{:next_event, :internal, :connect}]} 40 - end 41 - 42 - def disconnected(:enter, _from, data) do 43 - Logger.debug("Initial connection") 44 - # TODO: differentiate between initial & reconnects, probably stuff to do with seq 45 - {:next_state, :disconnected, data} 46 - end 25 + # Build opts for Drinkup.Socket from Options struct 26 + socket_opts = [ 27 + host: options.host, 28 + cursor: options.cursor, 29 + options: options 30 + ] 47 31 48 - def disconnected(:internal, :connect, data) do 49 - {:next_state, :connecting_http, data} 32 + Drinkup.Socket.start_link(__MODULE__, socket_opts, statem_opts) 50 33 end 51 34 52 - def connecting_http(:enter, _from, %{options: options} = data) do 53 - Logger.debug("Connecting to http") 54 - 55 - %{host: host, port: port} = URI.new!(options.host) 56 - 57 - {:ok, conn} = 58 - :gun.open(:binary.bin_to_list(host), port, %{ 59 - retry: 0, 60 - protocols: [:http], 61 - connect_timeout: @timeout, 62 - domain_lookup_timeout: @timeout, 63 - tls_handshake_timeout: @timeout, 64 - tls_opts: [ 65 - verify: :verify_peer, 66 - cacerts: :certifi.cacerts(), 67 - depth: 3, 68 - customize_hostname_check: [ 69 - match_fun: :public_key.pkix_verify_hostname_match_fun(:https) 70 - ] 71 - ] 72 - }) 73 - 74 - {:keep_state, %{data | conn: conn}, [{:state_timeout, @timeout, :connect_timeout}]} 35 + @impl true 36 + def build_path(%{seq: seq}) do 37 + cursor_param = if seq, do: %{cursor: seq}, else: %{} 38 + "/xrpc/com.atproto.sync.subscribeRepos?" <> URI.encode_query(cursor_param) 75 39 end 76 40 77 - def connecting_http(:info, {:gun_up, _conn, :http}, data) do 78 - {:next_state, :connecting_ws, data} 79 - end 80 - 81 - def connecting_http(:state_timeout, :connect_timeout, _data) do 82 - {:stop, :connect_http_timeout} 83 - end 84 - 85 - def connecting_ws(:enter, _from, %{conn: conn, seq: seq} = data) do 86 - Logger.debug("Upgrading connection to websocket") 87 - path = "/xrpc/com.atproto.sync.subscribeRepos?" <> URI.encode_query(%{cursor: seq}) 88 - stream = :gun.ws_upgrade(conn, path, [], %{flow: @flow}) 89 - {:keep_state, %{data | stream: stream}, [{:state_timeout, @timeout, :upgrade_timeout}]} 90 - end 91 - 92 - def connecting_ws(:info, {:gun_upgrade, _conn, _stream, ["websocket"], _headers}, data) do 93 - {:next_state, :connected, data} 94 - end 95 - 96 - def connecting_ws(:state_timeout, :upgrade_timeout, _data) do 97 - {:stop, :connect_ws_timeout} 98 - end 99 - 100 - def connected(:enter, _from, _data) do 101 - Logger.debug("Connected to websocket") 102 - :keep_state_and_data 103 - end 104 - 105 - def connected(:info, {:gun_ws, conn, stream, {:binary, frame}}, %{options: options} = data) do 106 - # TODO: let clients specify a handler for raw* (*decoded) packets to support any atproto subscription 107 - # Will also need support for JSON frames 41 + @impl true 42 + def handle_frame({:binary, frame}, %{seq: seq, options: options} = data) do 108 43 with {:ok, header, next} <- CAR.DagCbor.decode(frame), 109 44 {:ok, payload, _} <- CAR.DagCbor.decode(next), 110 45 {%{"op" => @op_regular, "t" => type}, _} <- {header, payload}, 111 - true <- Event.valid_seq?(data.seq, payload["seq"]) do 112 - data = %{data | seq: payload["seq"] || data.seq} 113 - message = Event.from(type, payload) 114 - :ok = :gun.update_flow(conn, stream, @flow) 46 + true <- Event.valid_seq?(seq, payload["seq"]) do 47 + new_seq = payload["seq"] || seq 115 48 116 - case message do 49 + case Event.from(type, payload) do 117 50 nil -> 118 51 Logger.warning("Received unrecognised event from firehose: #{inspect({type, payload})}") 119 52 ··· 121 54 Event.dispatch(message, options) 122 55 end 123 56 124 - {:keep_state, data} 57 + {:ok, %{data | seq: new_seq}} 125 58 else 126 59 false -> 127 60 Logger.error("Got out of sequence or invalid `seq` from Firehose") 128 - {:keep_state, data} 61 + :noop 129 62 130 63 {%{"op" => @op_error, "t" => type}, payload} -> 131 64 Logger.error("Got error from Firehose: #{inspect({type, payload})}") 132 - {:keep_state, data} 65 + :noop 133 66 134 67 {:error, reason} -> 135 68 Logger.warning("Failed to decode frame from Firehose: #{inspect(reason)}") 136 - {:keep_state, data} 69 + :noop 137 70 end 138 71 end 139 72 140 - def connected(:info, {:gun_ws, _conn, _stream, :close}, _data) do 73 + @impl true 74 + def handle_frame(:close, _data) do 141 75 Logger.info("Websocket closed, reason unknown") 142 - {:keep_state_and_data, [{:next_event, :internal, :reconnect}]} 76 + nil 143 77 end 144 78 145 - def connected(:info, {:gun_ws, _conn, _stream, {:close, errno, reason}}, _data) do 79 + @impl true 80 + def handle_frame({:close, errno, reason}, _data) do 146 81 Logger.info("Websocket closed, errno: #{errno}, reason: #{inspect(reason)}") 147 - {:keep_state_and_data, [{:next_event, :internal, :reconnect}]} 82 + nil 148 83 end 149 84 150 - def connected(:info, {:gun_down, old_conn, _proto, _reason, _killed_streams}, %{conn: new_conn}) 151 - when old_conn != new_conn do 152 - Logger.debug("Ignoring received :gun_down for a previous connection.") 153 - :keep_state_and_data 154 - end 155 - 156 - def connected(:info, {:gun_down, _conn, _proto, _reason, _killed_streams}, _data) do 157 - Logger.info("Websocket connection killed. Attempting to reconnect") 158 - {:keep_state_and_data, [{:next_event, :internal, :reconnect}]} 159 - end 160 - 161 - def connected(:internal, :reconnect, %{conn: conn} = data) do 162 - :ok = :gun.close(conn) 163 - :ok = :gun.flush(conn) 164 - 165 - # TODO: reconnect backoff 166 - {:next_state, :disconnected, %{data | conn: nil, stream: nil}, 167 - [{:next_event, :internal, :connect}]} 85 + @impl true 86 + def handle_frame({:text, _text}, _data) do 87 + Logger.warning("Received unexpected text frame from Firehose") 88 + :noop 168 89 end 169 90 end
+341
lib/socket.ex
··· 1 + defmodule Drinkup.Socket do 2 + # TODO: talk about how to implment, but that it's for internal use 3 + @moduledoc false 4 + 5 + require Logger 6 + 7 + @behaviour :gen_statem 8 + 9 + @type frame :: 10 + {:binary, binary()} 11 + | {:text, String.t()} 12 + | :close 13 + | {:close, errno :: integer(), reason :: binary()} 14 + 15 + @type user_data :: term() 16 + 17 + @type reconnect_strategy :: 18 + :exponential 19 + | {:exponential, max_backoff :: pos_integer()} 20 + | {:custom, (attempt :: pos_integer() -> delay_ms :: pos_integer())} 21 + 22 + @type option :: 23 + {:host, String.t()} 24 + | {:flow, pos_integer()} 25 + | {:timeout, pos_integer()} 26 + | {:tls_opts, keyword()} 27 + | {:gun_opts, map()} 28 + | {:reconnect_strategy, reconnect_strategy()} 29 + | {atom(), term()} 30 + 31 + @callback init(opts :: keyword()) :: {:ok, user_data()} | {:error, reason :: term()} 32 + 33 + @callback build_path(data :: user_data()) :: String.t() 34 + 35 + @callback handle_frame(frame :: frame(), data :: user_data()) :: 36 + {:ok, new_data :: user_data()} | :noop | nil | {:error, reason :: term()} 37 + 38 + @callback handle_connected(data :: user_data()) :: {:ok, new_data :: user_data()} 39 + 40 + @callback handle_disconnected(reason :: term(), data :: user_data()) :: 41 + {:ok, new_data :: user_data()} 42 + 43 + @optional_callbacks handle_connected: 1, handle_disconnected: 2 44 + 45 + defstruct [ 46 + :module, 47 + :user_data, 48 + :options, 49 + :conn, 50 + :stream, 51 + reconnect_attempts: 0 52 + ] 53 + 54 + defmacro __using__(_opts) do 55 + quote do 56 + @behaviour Drinkup.Socket 57 + 58 + def start_link(opts, statem_opts \\ []) 59 + 60 + def start_link(opts, statem_opts) do 61 + Drinkup.Socket.start_link(__MODULE__, opts, statem_opts) 62 + end 63 + 64 + defoverridable start_link: 2 65 + 66 + def child_spec(opts) do 67 + %{ 68 + id: __MODULE__, 69 + start: {__MODULE__, :start_link, [opts, []]}, 70 + type: :worker, 71 + restart: :permanent, 72 + shutdown: 500 73 + } 74 + end 75 + 76 + defoverridable child_spec: 1 77 + 78 + @impl true 79 + def handle_connected(data), do: {:ok, data} 80 + 81 + @impl true 82 + def handle_disconnected(_reason, data), do: {:ok, data} 83 + 84 + defoverridable handle_connected: 1, handle_disconnected: 2 85 + end 86 + end 87 + 88 + @impl true 89 + def callback_mode, do: [:state_functions, :state_enter] 90 + 91 + @doc """ 92 + Start a WebSocket connection process. 93 + 94 + ## Parameters 95 + 96 + * `module` - The module implementing the Drinkup.Socket behaviour 97 + * `opts` - Keyword list of options (see module documentation) 98 + * `statem_opts` - Options passed to `:gen_statem.start_link/3` 99 + """ 100 + def start_link(module, opts, statem_opts) do 101 + :gen_statem.start_link(__MODULE__, {module, opts}, statem_opts) 102 + end 103 + 104 + @impl true 105 + def init({module, opts}) do 106 + case module.init(opts) do 107 + {:ok, user_data} -> 108 + options = parse_options(opts) 109 + 110 + data = %__MODULE__{ 111 + module: module, 112 + user_data: user_data, 113 + options: options, 114 + reconnect_attempts: 0 115 + } 116 + 117 + {:ok, :disconnected, data, [{:next_event, :internal, :connect}]} 118 + 119 + {:error, reason} -> 120 + {:stop, {:init_failed, reason}} 121 + end 122 + end 123 + 124 + # :disconnected state - waiting to connect or reconnect 125 + 126 + def disconnected(:enter, _from, _data) do 127 + Logger.debug("[Drinkup.Socket] Entering disconnected state") 128 + :keep_state_and_data 129 + end 130 + 131 + def disconnected(:internal, :connect, data) do 132 + {:next_state, :connecting_http, data} 133 + end 134 + 135 + def disconnected(:timeout, :reconnect, data) do 136 + {:next_state, :connecting_http, data} 137 + end 138 + 139 + # :connecting_http state - establishing HTTP connection with TLS 140 + 141 + def connecting_http(:enter, _from, %{options: options} = data) do 142 + Logger.debug("[Drinkup.Socket] Connecting to HTTP") 143 + 144 + %{host: host, port: port} = URI.new!(options.host) 145 + 146 + gun_opts = 147 + Map.merge( 148 + %{ 149 + retry: 0, 150 + protocols: [:http], 151 + connect_timeout: options.timeout, 152 + domain_lookup_timeout: options.timeout, 153 + tls_handshake_timeout: options.timeout, 154 + tls_opts: options.tls_opts 155 + }, 156 + options.gun_opts 157 + ) 158 + 159 + case :gun.open(:binary.bin_to_list(host), port, gun_opts) do 160 + {:ok, conn} -> 161 + {:keep_state, %{data | conn: conn}, [{:state_timeout, options.timeout, :connect_timeout}]} 162 + 163 + {:error, reason} -> 164 + Logger.error("[Drinkup.Socket] Failed to open connection: #{inspect(reason)}") 165 + {:stop, {:connect_failed, reason}} 166 + end 167 + end 168 + 169 + def connecting_http(:info, {:gun_up, _conn, :http}, data) do 170 + {:next_state, :connecting_ws, data} 171 + end 172 + 173 + def connecting_http(:state_timeout, :connect_timeout, data) do 174 + Logger.error("[Drinkup.Socket] HTTP connection timeout") 175 + trigger_reconnect(data) 176 + end 177 + 178 + # :connecting_ws state - upgrading to WebSocket 179 + 180 + def connecting_ws( 181 + :enter, 182 + _from, 183 + %{module: module, user_data: user_data, options: options} = data 184 + ) do 185 + Logger.debug("[Drinkup.Socket] Upgrading connection to WebSocket") 186 + 187 + path = module.build_path(user_data) 188 + stream = :gun.ws_upgrade(data.conn, path, [], %{flow: options.flow}) 189 + 190 + {:keep_state, %{data | stream: stream}, [{:state_timeout, options.timeout, :upgrade_timeout}]} 191 + end 192 + 193 + def connecting_ws(:info, {:gun_upgrade, _conn, _stream, ["websocket"], _headers}, data) do 194 + {:next_state, :connected, data} 195 + end 196 + 197 + def connecting_ws(:info, {:gun_response, _conn, _stream, _fin, status, _headers}, data) do 198 + Logger.error("[Drinkup.Socket] WebSocket upgrade failed with status: #{status}") 199 + trigger_reconnect(data) 200 + end 201 + 202 + def connecting_ws(:info, {:gun_error, _conn, _stream, reason}, data) do 203 + Logger.error("[Drinkup.Socket] WebSocket upgrade error: #{inspect(reason)}") 204 + trigger_reconnect(data) 205 + end 206 + 207 + def connecting_ws(:state_timeout, :upgrade_timeout, data) do 208 + Logger.error("[Drinkup.Socket] WebSocket upgrade timeout") 209 + trigger_reconnect(data) 210 + end 211 + 212 + # :connected state - active WebSocket connection 213 + 214 + def connected(:enter, _from, %{module: module, user_data: user_data} = data) do 215 + Logger.debug("[Drinkup.Socket] WebSocket connected") 216 + 217 + case module.handle_connected(user_data) do 218 + {:ok, new_user_data} -> 219 + {:keep_state, %{data | user_data: new_user_data, reconnect_attempts: 0}} 220 + 221 + _ -> 222 + {:keep_state, %{data | reconnect_attempts: 0}} 223 + end 224 + end 225 + 226 + def connected( 227 + :info, 228 + {:gun_ws, conn, _stream, frame}, 229 + %{module: module, user_data: user_data, options: options} = data 230 + ) do 231 + result = module.handle_frame(frame, user_data) 232 + 233 + :ok = :gun.update_flow(conn, frame, options.flow) 234 + 235 + case result do 236 + {:ok, new_user_data} -> 237 + {:keep_state, %{data | user_data: new_user_data}} 238 + 239 + result when result in [:noop, nil] -> 240 + :keep_state_and_data 241 + 242 + {:error, reason} -> 243 + Logger.error("[Drinkup.Socket] Frame handler error: #{inspect(reason)}") 244 + :keep_state_and_data 245 + end 246 + end 247 + 248 + def connected(:info, {:gun_ws, _conn, _stream, :close}, data) do 249 + Logger.info("[Drinkup.Socket] WebSocket closed by remote") 250 + trigger_reconnect(data, :remote_close) 251 + end 252 + 253 + def connected(:info, {:gun_ws, _conn, _stream, {:close, errno, reason}}, data) do 254 + Logger.info("[Drinkup.Socket] WebSocket closed: #{errno} - #{inspect(reason)}") 255 + trigger_reconnect(data, {:remote_close, errno, reason}) 256 + end 257 + 258 + def connected(:info, {:gun_down, old_conn, _proto, _reason, _killed_streams}, %{conn: new_conn}) 259 + when old_conn != new_conn do 260 + Logger.debug("[Drinkup.Socket] Ignoring :gun_down for old connection") 261 + :keep_state_and_data 262 + end 263 + 264 + def connected(:info, {:gun_down, _conn, _proto, reason, _killed_streams}, data) do 265 + Logger.info("[Drinkup.Socket] Connection down: #{inspect(reason)}") 266 + trigger_reconnect(data, {:connection_down, reason}) 267 + end 268 + 269 + def connected( 270 + :internal, 271 + :reconnect, 272 + %{conn: conn, options: options, reconnect_attempts: attempts} = data 273 + ) do 274 + :ok = :gun.close(conn) 275 + :ok = :gun.flush(conn) 276 + 277 + backoff = calculate_backoff(attempts, options.reconnect_strategy) 278 + 279 + Logger.info("[Drinkup.Socket] Reconnecting in #{backoff}ms (attempt #{attempts + 1})") 280 + 281 + {:next_state, :disconnected, 282 + %{data | conn: nil, stream: nil, reconnect_attempts: attempts + 1}, 283 + [{{:timeout, :reconnect}, backoff, :reconnect}]} 284 + end 285 + 286 + # Helper functions 287 + 288 + defp trigger_reconnect(data, reason \\ :unknown) do 289 + %{module: module, user_data: user_data} = data 290 + 291 + case module.handle_disconnected(reason, user_data) do 292 + {:ok, new_user_data} -> 293 + {:keep_state, %{data | user_data: new_user_data}, [{:next_event, :internal, :reconnect}]} 294 + 295 + _ -> 296 + {:keep_state_and_data, [{:next_event, :internal, :reconnect}]} 297 + end 298 + end 299 + 300 + defp parse_options(opts) do 301 + %{ 302 + host: Keyword.fetch!(opts, :host), 303 + flow: Keyword.get(opts, :flow, 10), 304 + timeout: Keyword.get(opts, :timeout, :timer.seconds(5)), 305 + tls_opts: Keyword.get(opts, :tls_opts, default_tls_opts()), 306 + gun_opts: Keyword.get(opts, :gun_opts, %{}), 307 + reconnect_strategy: Keyword.get(opts, :reconnect_strategy, :exponential) 308 + } 309 + end 310 + 311 + defp default_tls_opts do 312 + [ 313 + verify: :verify_peer, 314 + cacerts: :certifi.cacerts(), 315 + depth: 3, 316 + customize_hostname_check: [ 317 + match_fun: :public_key.pkix_verify_hostname_match_fun(:https) 318 + ] 319 + ] 320 + end 321 + 322 + defp calculate_backoff(attempt, strategy) do 323 + case strategy do 324 + :exponential -> 325 + exponential_backoff(attempt, :timer.seconds(60)) 326 + 327 + {:exponential, max_backoff} -> 328 + exponential_backoff(attempt, max_backoff) 329 + 330 + {:custom, func} when is_function(func, 1) -> 331 + func.(attempt) 332 + end 333 + end 334 + 335 + defp exponential_backoff(attempt, max_backoff) do 336 + base = :timer.seconds(1) 337 + delay = min(base * :math.pow(2, attempt), max_backoff) 338 + jitter = :rand.uniform(trunc(delay * 0.1)) 339 + trunc(delay) + jitter 340 + end 341 + end