Elixir ATProtocol firehose & subscription listener

refactor: rewrite socket using gen_statem and gun

Inspired by Nostrum's sturcture.

ovyerus.com 7c948f8f 4ec6c6d7

verified
+1 -85
lib/drinkup.ex
···
-
defmodule Drinkup.Firehose do
-
alias Drinkup.Firehose
-
require Logger
-
-
use WebSockex
-
-
@default_host "https://bsky.network"
-
-
@op_regular 1
-
@op_error -1
-
-
# TODO: switch to Gun and GenServer?
-
-
def start_link(opts \\ []) do
-
opts = Keyword.validate!(opts, host: @default_host)
-
host = Keyword.get(opts, :host)
-
cursor = Keyword.get(opts, :cursor)
-
-
url =
-
"#{host}/xrpc/com.atproto.sync.subscribeRepos"
-
|> URI.new!()
-
|> URI.append_query(URI.encode_query(%{cursor: cursor}))
-
|> URI.to_string()
-
-
WebSockex.start_link(url, __MODULE__, 0)
-
end
-
-
def handle_connect(conn, state) do
-
Logger.info("Connected to Firehose at #{conn.host}#{conn.path}")
-
{:ok, state}
-
end
-
-
def handle_frame({:binary, msg}, state) do
-
with {:ok, header, next} <- CAR.DagCbor.decode(msg),
-
{:ok, payload, _} <- CAR.DagCbor.decode(next),
-
{%{"op" => @op_regular, "t" => type}, _} <- {header, payload},
-
true <- type == "#info" || valid_seq?(state, payload["seq"]),
-
message <-
-
from_payload(type, payload) do
-
case message do
-
%Firehose.Commit{} = commit ->
-
IO.inspect(commit.ops, label: commit.repo)
-
-
msg ->
-
IO.inspect(msg)
-
end
-
-
{:ok, payload["seq"] || state}
-
else
-
false ->
-
Logger.error("Got out of sequence or invalid `seq` from Firehose")
-
{:ok, state}
-
-
{%{"op" => @op_error, "t" => type}, payload} ->
-
Logger.error("Got error from Firehose: #{inspect({type, payload})}")
-
{:ok, state}
-
-
{:error, reason} ->
-
Logger.warning("Failed to decode frame from Firehose: #{inspect(reason)}")
-
{:ok, state}
-
end
-
end
-
-
def handle_frame({:text, msg}, state) do
-
Logger.warning("Got unexpected text frame from Firehose: #{inspect(msg)}")
-
{:ok, state}
-
end
-
-
@spec valid_seq?(integer(), any()) :: boolean()
-
defp valid_seq?(last_seq, seq) when is_integer(seq), do: seq > last_seq
-
defp valid_seq?(_last_seq, _seq), do: false
-
-
@spec from_payload(String.t(), map()) ::
-
Firehose.Commit.t()
-
| Firehose.Sync.t()
-
| Firehose.Identity.t()
-
| Firehose.Account.t()
-
| Firehose.Info.t()
-
| nil
-
defp from_payload("#commit", payload), do: Firehose.Commit.from(payload)
-
defp from_payload("#sync", payload), do: Firehose.Sync.from(payload)
-
defp from_payload("#identity", payload), do: Firehose.Identity.from(payload)
-
defp from_payload("#account", payload), do: Firehose.Account.from(payload)
-
defp from_payload("#info", payload), do: Firehose.Info.from(payload)
-
defp from_payload(_type, _payload), do: nil
+
defmodule Drinkup do
end
+22
lib/event.ex
···
+
defmodule Drinkup.Event do
+
alias Drinkup.Event
+
+
@spec from(String.t(), map()) ::
+
Event.Commit.t()
+
| Event.Sync.t()
+
| Event.Identity.t()
+
| Event.Account.t()
+
| Event.Info.t()
+
| nil
+
def from("#commit", payload), do: Event.Commit.from(payload)
+
def from("#sync", payload), do: Event.Sync.from(payload)
+
def from("#identity", payload), do: Event.Identity.from(payload)
+
def from("#account", payload), do: Event.Account.from(payload)
+
def from("#info", payload), do: Event.Info.from(payload)
+
def from(_type, _payload), do: nil
+
+
@spec valid_seq?(integer() | nil, any()) :: boolean()
+
def valid_seq?(nil, seq) when is_integer(seq), do: true
+
def valid_seq?(last_seq, seq) when is_integer(last_seq) and is_integer(seq), do: seq > last_seq
+
def valid_seq?(_last_seq, _seq), do: false
+
end
+1 -1
lib/firehose/account.ex lib/event/account.ex
···
-
defmodule Drinkup.Firehose.Account do
+
defmodule Drinkup.Event.Account do
@moduledoc """
Struct for account events from the ATProto Firehose.
"""
+1 -1
lib/firehose/commit.ex lib/event/commit.ex
···
-
defmodule Drinkup.Firehose.Commit do
+
defmodule Drinkup.Event.Commit do
@moduledoc """
Struct for commit events from the ATProto Firehose.
"""
+1 -1
lib/firehose/identity.ex lib/event/identity.ex
···
-
defmodule Drinkup.Firehose.Identity do
+
defmodule Drinkup.Event.Identity do
@moduledoc """
Struct for identity events from the ATProto Firehose.
"""
+1 -1
lib/firehose/info.ex lib/event/info.ex
···
-
defmodule Drinkup.Firehose.Info do
+
defmodule Drinkup.Event.Info do
@moduledoc """
Struct for info events from the ATProto Firehose.
"""
+1 -1
lib/firehose/sync.ex lib/event/sync.ex
···
-
defmodule Drinkup.Firehose.Sync do
+
defmodule Drinkup.Event.Sync do
@moduledoc """
Struct for sync events from the ATProto Firehose.
"""
+165
lib/socket.ex
···
+
defmodule Drinkup.Socket do
+
@moduledoc """
+
gen_statem process for managing the websocket connection to an ATProto relay.
+
"""
+
+
require Logger
+
alias Drinkup.Event
+
+
@behaviour :gen_statem
+
@default_host "https://bsky.network"
+
@timeout :timer.seconds(5)
+
# TODO: `flow` determines messages in buffer. Determine ideal value?
+
@flow 10
+
+
@op_regular 1
+
@op_error -1
+
+
defstruct [:host, :seq, :conn, :stream]
+
+
@impl true
+
def callback_mode, do: [:state_functions, :state_enter]
+
+
def start_link(opts \\ []) do
+
opts = Keyword.validate!(opts, host: @default_host)
+
host = Keyword.get(opts, :host)
+
cursor = Keyword.get(opts, :cursor)
+
+
:gen_statem.start_link(__MODULE__, {host, cursor}, [])
+
end
+
+
@impl true
+
def init({host, cursor}) do
+
data = %__MODULE__{host: host, seq: cursor}
+
{:ok, :disconnected, data, [{:next_event, :internal, :connect}]}
+
end
+
+
def disconnected(:enter, _from, data) do
+
Logger.debug("Initial connection")
+
# TODO: differentiate between initial & reconnects, probably stuff to do with seq
+
{:next_state, :disconnected, data}
+
end
+
+
def disconnected(:internal, :connect, data) do
+
{:next_state, :connecting_http, data}
+
end
+
+
def connecting_http(:enter, _from, data) do
+
Logger.debug("Connecting to http")
+
+
%{host: host, port: port} = URI.new!(data.host)
+
+
{:ok, conn} =
+
:gun.open(:binary.bin_to_list(host), port, %{
+
retry: 0,
+
protocols: [:http],
+
connect_timeout: @timeout,
+
domain_lookup_timeout: @timeout,
+
tls_handshake_timeout: @timeout,
+
tls_opts: [
+
verify: :verify_peer,
+
cacerts: :certifi.cacerts(),
+
depth: 3,
+
customize_hostname_check: [
+
match_fun: :public_key.pkix_verify_hostname_match_fun(:https)
+
]
+
]
+
})
+
+
{:keep_state, %{data | conn: conn}, [{:state_timeout, @timeout, :connect_timeout}]}
+
end
+
+
def connecting_http(:info, {:gun_up, _conn, :http}, data) do
+
{:next_state, :connecting_ws, data}
+
end
+
+
def connecting_http(:state_timeout, :connect_timeout, _data) do
+
{:stop, :connect_http_timeout}
+
end
+
+
def connecting_ws(:enter, _from, %{conn: conn, seq: seq} = data) do
+
Logger.debug("Upgrading connection to websocket")
+
path = "/xrpc/com.atproto.sync.subscribeRepos?" <> URI.encode_query(%{cursor: seq})
+
stream = :gun.ws_upgrade(conn, path, [], %{flow: @flow})
+
{:keep_state, %{data | stream: stream}, [{:state_timeout, @timeout, :upgrade_timeout}]}
+
end
+
+
def connecting_ws(:info, {:gun_upgrade, _conn, _stream, ["websocket"], _headers}, data) do
+
{:next_state, :connected, data}
+
end
+
+
def connecting_ws(:state_timeout, :upgrade_timeout, _data) do
+
{:stop, :connect_ws_timeout}
+
end
+
+
def connected(:enter, _from, _data) do
+
Logger.debug("Connected to websocket")
+
:keep_state_and_data
+
end
+
+
def connected(:info, {:gun_ws, conn, stream, {:binary, frame}}, data) do
+
# TODO: let clients specify a handler for raw* (*decoded) packets to support any atproto subscription
+
# Will also need support for JSON frames
+
with {:ok, header, next} <- CAR.DagCbor.decode(frame),
+
{:ok, payload, _} <- CAR.DagCbor.decode(next),
+
{%{"op" => @op_regular, "t" => type}, _} <- {header, payload},
+
true <- type == "#info" || Event.valid_seq?(data.seq, payload["seq"]),
+
data <- %{data | seq: payload["seq"] || data.seq},
+
message <-
+
Event.from(type, payload) do
+
:ok = :gun.update_flow(conn, stream, @flow)
+
+
case message do
+
%Event.Commit{} = commit ->
+
IO.inspect(commit.ops, label: commit.repo)
+
+
msg ->
+
IO.inspect(msg)
+
end
+
+
{:keep_state, data}
+
else
+
false ->
+
Logger.error("Got out of sequence or invalid `seq` from Firehose")
+
{:keep_state, data}
+
+
{%{"op" => @op_error, "t" => type}, payload} ->
+
Logger.error("Got error from Firehose: #{inspect({type, payload})}")
+
{:keep_state, data}
+
+
{:error, reason} ->
+
Logger.warning("Failed to decode frame from Firehose: #{inspect(reason)}")
+
{:keep_state, data}
+
end
+
end
+
+
def connected(:info, {:gun_ws, _conn, _stream, :close}, _data) do
+
Logger.info("Websocket closed, reason unknown")
+
{:keep_state_and_data, [{:next_event, :internal, :reconnect}]}
+
end
+
+
def connected(:info, {:gun_ws, _conn, _stream, {:close, errno, reason}}, _data) do
+
Logger.info("Websocket closed, errno: #{errno}, reason: #{inspect(reason)}")
+
{:keep_state_and_data, [{:next_event, :internal, :reconnect}]}
+
end
+
+
def connected(:info, {:gun_down, old_conn, _proto, _reason, _killed_streams}, %{conn: new_conn})
+
when old_conn != new_conn do
+
Logger.debug("Ignoring received :gun_down for a previous connection.")
+
:keep_state_and_data
+
end
+
+
def connected(:info, {:gun_down, _conn, _proto, _reason, _killed_streams}, _data) do
+
Logger.info("Websocket connection killed. Attempting to reconnect")
+
{:keep_state_and_data, [{:next_event, :internal, :reconnect}]}
+
end
+
+
def connected(:internal, :reconnect, %{conn: conn} = data) do
+
:ok = :gun.close(conn)
+
:ok = :gun.flush(conn)
+
+
# TODO: reconnect backoff
+
{:next_state, :disconnected, %{data | conn: nil, stream: nil},
+
[{:next_event, :internal, :connect}]}
+
end
+
end
+4 -3
mix.exs
···
# Run "mix help deps" to learn about dependencies.
defp deps do
[
-
{:cbor, "~> 1.0.0"},
{:car, "~> 0.1.0"},
+
{:cbor, "~> 1.0.0"},
+
{:certifi, "~> 2.15"},
{:credo, "~> 1.7", only: [:dev, :test], runtime: false},
-
{:typedstruct, "~> 0.5"},
-
{:websockex, "~> 0.5.0", hex: :websockex_wt}
+
{:gun, "~> 2.2"},
+
{:typedstruct, "~> 0.5"}
]
end
end
+3 -2
mix.lock
···
"bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
"car": {:hex, :car, "0.1.1", "a5bc4c5c1be96eab437634b3c0ccad1fe17b5e3d68c22a4031241ae1345aebd4", [:mix], [{:cbor, "~> 1.0.0", [hex: :cbor, repo: "hexpm", optional: false]}, {:typedstruct, "~> 0.5", [hex: :typedstruct, repo: "hexpm", optional: false]}, {:varint, "~> 1.4", [hex: :varint, repo: "hexpm", optional: false]}], "hexpm", "f895dda8123d04dd336db5a2bf0d0b47f4559cd5383f83fcca0700c1b45bfb6a"},
"cbor": {:hex, :cbor, "1.0.1", "39511158e8ea5a57c1fcb9639aaa7efde67129678fee49ebbda780f6f24959b0", [:mix], [], "hexpm", "5431acbe7a7908f17f6a9cd43311002836a34a8ab01876918d8cfb709cd8b6a2"},
+
"certifi": {:hex, :certifi, "2.15.0", "0e6e882fcdaaa0a5a9f2b3db55b1394dba07e8d6d9bcad08318fb604c6839712", [:rebar3], [], "hexpm", "b147ed22ce71d72eafdad94f055165c1c182f61a2ff49df28bcc71d1d5b94a60"},
+
"cowlib": {:hex, :cowlib, "2.15.0", "3c97a318a933962d1c12b96ab7c1d728267d2c523c25a5b57b0f93392b6e9e25", [:make, :rebar3], [], "hexpm", "4f00c879a64b4fe7c8fcb42a4281925e9ffdb928820b03c3ad325a617e857532"},
"credo": {:hex, :credo, "1.7.12", "9e3c20463de4b5f3f23721527fcaf16722ec815e70ff6c60b86412c695d426c1", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "8493d45c656c5427d9c729235b99d498bd133421f3e0a683e5c1b561471291e5"},
"file_system": {:hex, :file_system, "1.1.0", "08d232062284546c6c34426997dd7ef6ec9f8bbd090eb91780283c9016840e8f", [:mix], [], "hexpm", "bfcf81244f416871f2a2e15c1b515287faa5db9c6bcf290222206d120b3d43f6"},
+
"gun": {:hex, :gun, "2.2.0", "b8f6b7d417e277d4c2b0dc3c07dfdf892447b087f1cc1caff9c0f556b884e33d", [:make, :rebar3], [{:cowlib, ">= 2.15.0 and < 3.0.0", [hex: :cowlib, repo: "hexpm", optional: false]}], "hexpm", "76022700c64287feb4df93a1795cff6741b83fb37415c40c34c38d2a4645261a"},
"jason": {:hex, :jason, "1.4.4", "b9226785a9aa77b6857ca22832cffa5d5011a667207eb2a0ad56adb5db443b8a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "c5eb0cab91f094599f94d55bc63409236a8ec69a21a67814529e8d5f6cc90b3b"},
-
"telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"},
"typedstruct": {:hex, :typedstruct, "0.5.3", "d68ae424251a41b81a8d0c485328ab48edbd3858f3565bbdac21b43c056fc9b4", [:make, :mix], [], "hexpm", "b53b8186701417c0b2782bf02a2db5524f879b8488f91d1d83b97d84c2943432"},
"varint": {:hex, :varint, "1.5.1", "17160c70d0428c3f8a7585e182468cac10bbf165c2360cf2328aaa39d3fb1795", [:mix], [], "hexpm", "24f3deb61e91cb988056de79d06f01161dd01be5e0acae61d8d936a552f1be73"},
-
"websockex": {:hex, :websockex_wt, "0.5.0", "8725b3bc741e7a682c21310610d033f0aaeedfb3238d9c8d5522c345c04f3f93", [:mix], [{:telemetry, "~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "f854a5e7dbd61e852ee74565d862d606e884425c8867ddaa49c8a6c6f43d832a"},
}