defmodule Nats.Streaming.Client do @behaviour :gen_statem @enforce_keys [:client_id, :conn_id, :connection_name] defstruct client_id: nil, cluster_id: "test-cluster", conn_id: nil, connection_name: nil, connection_pid: nil, close_subject: nil, heartbeat_subject: nil, pub_subject: nil, sub_subject: nil, unsub_subject: nil @type t :: %__MODULE__{ client_id: String.t(), cluster_id: String.t(), conn_id: String.t(), connection_name: atom(), connection_pid: pid() | nil, close_subject: String.t() | nil, pub_subject: String.t() | nil, sub_subject: String.t() | nil, unsub_subject: String.t() | nil } require Logger alias Nats.Streaming.Protocol def start_link(settings, options \\ []) do case Keyword.get(options, :name) do nil -> :gen_statem.start_link(__MODULE__, settings, options) name when is_atom(name) -> :gen_statem.start_link({:local, name}, __MODULE__, settings, options) name -> :gen_statem.start_link(name, __MODULE__, settings, options) end end @spec pub(GenServer.server(), String.t(), binary(), keyword()) :: :ok | {:error, term()} def pub(streaming_client, subject, payload, options \\ []) do reply_to = Keyword.get(options, :reply_to) guid = Keyword.get(options, :guid) || nuid() with {:ok, {client_id, pub_subject, connection_pid}} <- pub_info(streaming_client), pub_msg <- encode_pub_msg(subject, payload, guid, reply_to, client_id), {:ok, %{body: pb}} <- Gnat.request(connection_pid, pub_subject, pub_msg) do case Protocol.PubAck.decode(pb) do %Protocol.PubAck{error: ""} -> :ok %Protocol.PubAck{error: err} -> {:error, err} end end end @spec pub_info(GenServer.server()) :: {:ok, {String.t(), String.t(), pid()}} | {:error, term()} def pub_info(streaming_client) do :gen_statem.call(streaming_client, :pub_info) end @spec sub_info(GenServer.server()) :: {:ok, {String.t(), String.t(), pid()}} | {:error, term()} def sub_info(streaming_client) do :gen_statem.call(streaming_client, :sub_info) end # Callback Functions @impl :gen_statem def callback_mode(), do: :state_functions @impl :gen_statem def init(settings) do Process.flag(:trap_exit, true) state = new(settings) {:ok, :disconnected, state, [{:next_event, :internal, :connect}]} end @impl :gen_statem def terminate(:shutdown, state, _data) do req = Protocol.CloseRequest.new(clientID: state.client_id) |> Protocol.CloseRequest.encode() case Gnat.request(state.connection_pid, state.close_subject, req) do {:ok, %{body: msg}} -> msg = Protocol.CloseResponse.decode(msg) Logger.debug("#{__MODULE__} TODO - Handle error attribute #{inspect(msg)}") {:error, reason} -> Logger.error("Failed to close connection: #{inspect(reason)}") end :ok end def terminate(reason, _state, _data) do Logger.error("#{__MODULE__} unexpected shutdown #{inspect(reason)}") end # Internal State Functions @doc false def new(settings) do connection_name = Keyword.fetch!(settings, :connection_name) client_id = Keyword.get(settings, :client_id) || :crypto.strong_rand_bytes(12) |> Base.encode16() conn_id = Keyword.get(settings, :conn_id) || :crypto.strong_rand_bytes(12) |> Base.encode16() heartbeat_subject = "#{client_id}.#{conn_id}.heartbeat" %__MODULE__{ client_id: client_id, cluster_id: Keyword.get(settings, :cluster_id, "test-cluster"), conn_id: conn_id, connection_name: connection_name, heartbeat_subject: heartbeat_subject } end @doc false def disconnected(:internal, :connect, %__MODULE__{connection_name: connection_name}) do maybe_pid = Process.whereis(connection_name) {:keep_state_and_data, [{:next_event, :internal, {:find_connection, maybe_pid}}]} end def disconnected({:timeout, :reconnect}, _, state), do: disconnected(:internal, :connect, state) def disconnected(:internal, {:find_connection, nil}, _state) do {:keep_state_and_data, [{{:timeout, :reconnect}, 250, :reconnect}]} end def disconnected(:internal, {:find_connection, pid}, state) when is_pid(pid) do state = %__MODULE__{state | connection_pid: pid} actions = [{:next_event, :internal, :monitor_and_subscribe}] {:next_state, :connected, state, actions} end def disconnected({:call, from}, _call, _state) do {:keep_state_and_data, [{:reply, from, {:error, :disconnected}}]} end @doc false def connected(:internal, :monitor_and_subscribe, %__MODULE__{} = state) do _ref = Process.monitor(state.connection_pid) {:ok, _sid} = Gnat.sub(state.connection_pid, self(), state.heartbeat_subject) actions = [{:next_event, :internal, :register}] {:keep_state_and_data, actions} end def connected(:internal, :register, %__MODULE__{} = state) do req = Protocol.ConnectRequest.new( clientID: state.client_id, connID: state.conn_id, heartbeatInbox: state.heartbeat_subject ) |> Protocol.ConnectRequest.encode() case Gnat.request(state.connection_pid, "_STAN.discover.#{state.cluster_id}", req) do {:ok, %{body: msg}} -> msg = Protocol.ConnectResponse.decode(msg) actions = [{:next_event, :internal, {:connect_response, msg}}] {:keep_state_and_data, actions} {:error, reason} -> Logger.error("Failed to connect to NATS Streaming server: #{inspect(reason)}") actions = [{{:timeout, :reregister}, 1_000, :reregister}] {:keep_state_and_data, actions} end end def connected({:timeout, :reregister}, _, state), do: connected(:internal, :register, state) def connected(:internal, {:connect_response, response}, %__MODULE__{} = state) do if response.error == "" do state = state |> Map.put(:close_subject, response.closeRequests) |> Map.put(:pub_subject, "#{response.pubPrefix}.#{state.client_id}.#{state.conn_id}") |> Map.put(:sub_subject, response.subRequests) |> Map.put(:unsub_subject, response.unsubRequests) {:next_state, :registered, state, []} else Logger.error("Got Connection Error From NATS Streaming server #{response.error}") {:keep_state_and_data, [{{:timeout, :reregister}, 1_000, :reregister}]} end end def connected( :info, {:DOWN, _ref, :process, pid, _reason}, %__MODULE__{connection_pid: pid} = state ) do state = %__MODULE__{state | connection_pid: nil} {:next_state, :disconnected, state, [{{:timeout, :reconnect}, 250, :reconnect}]} end def connected({:call, from}, _call, _state) do {:keep_state_and_data, [{:reply, from, {:error, :not_registered}}]} end @doc false def registered({:call, from}, :pub_info, %__MODULE__{} = state) do pub_info = {state.client_id, state.pub_subject, state.connection_pid} {:keep_state_and_data, [{:reply, from, {:ok, pub_info}}]} end def registered({:call, from}, :sub_info, %__MODULE__{} = state) do sub_info = {state.client_id, state.sub_subject, state.connection_pid} {:keep_state_and_data, [{:reply, from, {:ok, sub_info}}]} end def registered(:info, {:msg, %{body: "", reply_to: reply_to}}, _state) do {:keep_state_and_data, [{:next_event, :internal, {:pub, reply_to, ""}}]} end def registered(:internal, {:pub, subject, body}, %__MODULE__{} = state) do :ok = Gnat.pub(state.connection_pid, subject, body) {:keep_state_and_data, []} end def registered( :info, {:DOWN, _ref, :process, pid, _reason}, %__MODULE__{connection_pid: pid} = state ) do state = %__MODULE__{ state | connection_pid: nil, close_subject: nil, pub_subject: nil, sub_subject: nil, unsub_subject: nil } {:next_state, :disconnected, state, [{{:timeout, :reconnect}, 250, :reconnect}]} end defp encode_pub_msg(subject, payload, guid, reply_to, client_id) do Protocol.PubMsg.new( clientID: client_id, guid: guid, subject: subject, data: payload, reply: reply_to ) |> Protocol.PubMsg.encode() end defp nuid, do: :crypto.strong_rand_bytes(12) |> Base.encode64() end