Something went wrong. Try again.
Elixir ATProtocol ingestion and sync library.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202defmodule Drinkup.Jetstream.Socket do @moduledoc """ WebSocket connection handler for Jetstream event streams.
Implements the Drinkup.Socket behaviour to manage connections to a Jetstream service, handling zstd-compressed JSON events and dispatching them to the configured consumer. """
use Drinkup.Socket
require Logger alias Drinkup.Jetstream.{Event, Options}
@dict_path "priv/jetstream/zstd_dictionary" @external_resource @dict_path @zstd_dict File.read!(@dict_path)
@impl true def init(opts) do options = Keyword.fetch!(opts, :options)
{:ok, %{options: options, host: options.host, cursor: options.cursor}} end
def start_link(%Options{} = options, statem_opts) do socket_opts = [ host: options.host, options: options ]
statem_opts = Keyword.put( statem_opts, :name, {:via, Registry, {Drinkup.Registry, {options.name, JetstreamSocket}}} )
Drinkup.Socket.start_link(__MODULE__, socket_opts, statem_opts) end
@impl true def build_path(%{options: options}) do query_params = [compress: "true"]
query_params = query_params |> put_collections(options.wanted_collections) |> put_dids(options.wanted_dids) |> put_cursor(options.cursor) |> put_max_size(options.max_message_size_bytes) |> put_require_hello(options.require_hello)
"/subscribe?" <> URI.encode_query(query_params) end
@impl true def handle_frame( {:binary, compressed_data}, {%{options: options} = data, _conn, _stream} ) do case decompress_and_parse(compressed_data) do {:ok, payload} -> case Event.from(payload) do nil -> # Event.from already logs warnings for unrecognized events :noop
event -> Event.dispatch(event, options) # Update cursor with the event's time_us new_cursor = Map.get(payload, "time_us") {:ok, %{data | cursor: new_cursor}} end
# TODO: sometimes getting ZSTD_CONTENTSIZE_UNKNOWN {:error, reason} -> Logger.error( "[Drinkup.Jetstream.Socket] Failed to decompress/parse frame: #{inspect(reason)}" )
:noop end end
@impl true def handle_frame({:text, json}, {%{options: options} = data, _conn, _stream}) do # Text frames shouldn't happen since we force compression, but handle them anyway case Jason.decode(json) do {:ok, payload} -> case Event.from(payload) do nil -> :noop
event -> Event.dispatch(event, options) new_cursor = Map.get(payload, "time_us") {:ok, %{data | cursor: new_cursor}} end
{:error, reason} -> Logger.error("[Drinkup.Jetstream.Socket] Failed to decode JSON: #{inspect(reason)}") :noop end end
@impl true def handle_frame(:close, _data) do Logger.info("[Drinkup.Jetstream.Socket] WebSocket closed, reason unknown") nil end
@impl true def handle_frame({:close, errno, reason}, _data) do Logger.info( "[Drinkup.Jetstream.Socket] WebSocket closed, errno: #{errno}, reason: #{inspect(reason)}" )
nil end
@impl true def handle_connected({user_data, conn, stream}) do # Register connection for options updates Registry.register( Drinkup.Registry, {user_data.options.name, JetstreamConnection}, {conn, stream} )
{:ok, user_data} end
@impl true def handle_disconnected(_reason, {user_data, _conn, _stream}) do # Unregister connection when disconnected Registry.unregister(Drinkup.Registry, {user_data.options.name, JetstreamConnection}) {:ok, user_data} end
# Can't use `create_ddict` as the value of `@zstd_dict` because it returns a reference :( @spec get_dictionary() :: reference() defp get_dictionary() do case :ezstd.create_ddict(@zstd_dict) do {:error, reason} -> raise ArgumentError, "somehow failed to created Jetstream's ZSTD dictionary: #{inspect(reason)}"
dict -> dict end end
@spec decompress_and_parse(binary()) :: {:ok, map()} | {:error, term()} defp decompress_and_parse(compressed_data) do with ctx when is_reference(ctx) <- :ezstd.create_decompression_context(byte_size(compressed_data)), :ok <- :ezstd.select_ddict(ctx, get_dictionary()), iolist when is_list(iolist) <- :ezstd.decompress_streaming(ctx, compressed_data), decompressed <- IO.iodata_to_binary(iolist), {:ok, payload} <- JSON.decode(decompressed) do {:ok, payload} else {:error, reason} -> {:error, reason} end end
@spec put_collections(keyword(), [String.t()]) :: keyword() defp put_collections(params, []), do: params
defp put_collections(params, collections) when is_list(collections) do Enum.reduce(collections, params, fn collection, acc -> [{:wantedCollections, collection} | acc] end) end
@spec put_dids(keyword(), [String.t()]) :: keyword() defp put_dids(params, []), do: params
defp put_dids(params, dids) when is_list(dids) do Enum.reduce(dids, params, fn did, acc -> [{:wantedDids, did} | acc] end) end
@spec put_cursor(keyword(), integer() | nil) :: keyword() defp put_cursor(params, nil), do: params
defp put_cursor(params, cursor) when is_integer(cursor), do: [{:cursor, cursor} | params]
@spec put_max_size(keyword(), integer() | nil) :: keyword() defp put_max_size(params, nil), do: params
defp put_max_size(params, max_size) when is_integer(max_size), do: [{:maxMessageSizeBytes, max_size} | params]
@spec put_require_hello(keyword(), boolean()) :: keyword() defp put_require_hello(params, false), do: params
defp put_require_hello(params, true), do: [{:requireHello, "true"} | params]end