defmodule NSQ.Connection do
@moduledoc """
Sets up a TCP connection to NSQD. Both consumers and producers use this.
"""
# ------------------------------------------------------- #
# Directives #
# ------------------------------------------------------- #
use GenServer
alias NSQ.Connection.Command
alias NSQ.Connection.Initializer
alias NSQ.Connection.MessageHandling
alias NSQ.ConnInfo
# ------------------------------------------------------- #
# Type Definitions #
# ------------------------------------------------------- #
@typedoc """
A tuple with a host and a port.
"""
@type host_with_port :: {String.t(), integer}
@typedoc """
A tuple with a string ID (used to target the connection in
NSQ.Connection.Supervisor) and a PID of the connection.
"""
@type connection :: {String.t(), pid}
@typedoc """
A map, but we can be more specific by asserting some entries that should be
set for a connection's state map.
"""
@type state :: %{parent: pid, config: NSQ.Config.t(), nsqd: host_with_port}
# ------------------------------------------------------- #
# Module Attributes #
# ------------------------------------------------------- #
@initial_state %{
parent: nil,
socket: nil,
reader: nil,
writer: nil,
connected: nil,
cmd_resp_queue: :queue.new(),
cmd_queue: :queue.new(),
config: %{},
reader_pid: nil,
msg_sup_pid: nil,
event_manager_pid: nil,
messages_in_flight: 0,
nsqd: nil,
topic: nil,
channel: nil,
backoff_counter: 0,
backoff_duration: 0,
max_rdy: 2500,
connect_attempts: 0,
stop_flag: false,
conn_info_pid: nil,
msg_timeout: nil
}
# ------------------------------------------------------- #
# Behaviour Implementation #
# ------------------------------------------------------- #
@spec start_link({pid | host_with_port | NSQ.Config.t() | String.t() | String.t() | pid}) ::
{:ok, pid}
def start_link({parent, nsqd, config, topic, channel, conn_info_pid, event_manager_pid}),
do: start_link({parent, nsqd, config, topic, channel, conn_info_pid, event_manager_pid, []})
@spec start_link({pid | host_with_port | NSQ.Config.t() | String.t() | String.t() | pid | list}) ::
{:ok, pid}
def start_link({parent, nsqd, config, topic, channel, conn_info_pid, event_manager_pid, opts}) do
state = %{
@initial_state
| parent: parent,
nsqd: nsqd,
config: config,
topic: topic,
channel: channel,
conn_info_pid: conn_info_pid,
event_manager_pid: event_manager_pid
}
{:ok, _pid} = GenServer.start_link(__MODULE__, state, opts)
end
@spec init(state) :: {:ok, state}
def init(conn_state) do
{:ok, reader} = NSQ.Connection.Buffer.start_link(:reader)
conn_state = %{conn_state | reader: reader}
{:ok, writer} = NSQ.Connection.Buffer.start_link(:writer)
conn_state = %{conn_state | writer: writer}
{:ok, msg_sup_pid} = NSQ.Message.Supervisor.start_link()
conn_state = %{conn_state | msg_sup_pid: msg_sup_pid}
conn_state |> ConnInfo.init()
case conn_state |> Initializer.connect() do
{:ok, state} -> {:ok, state}
{{:error, _reason}, state} -> {:ok, state}
end
end
def terminate(_reason, _state) do
:ok
end
@spec handle_call({:cmd, tuple, atom}, {pid, reference}, state) ::
{:reply, {:ok, reference}, state}
| {:reply, {:queued, :nosocket}, state}
def handle_call({:cmd, cmd, kind}, from, state) do
{reply, state} = state |> Command.exec(cmd, kind, from)
{:reply, reply, state}
end
@spec handle_call(:stop, {pid, reference}, state) :: {:stop, :normal, state}
def handle_call(:stop, _from, state) do
{:stop, :normal, state}
end
@spec handle_call(:state, {pid, reference}, state) :: {:reply, state, state}
def handle_call(:state, _from, state) do
{:reply, state, state}
end
@spec handle_call({:nsq_msg, binary}, {pid, reference}, state) :: {:reply, :ok, state}
def handle_call({:nsq_msg, msg}, _from, state) do
{:ok, state} = MessageHandling.handle_nsq_message(msg, state)
{:reply, :ok, state}
end
@spec handle_cast(:flush_cmd_queue, state) :: {:noreply, state}
def handle_cast(:flush_cmd_queue, state) do
{:noreply, Command.flush_cmd_queue!(state)}
end
@spec handle_cast(:reconnect, state) :: {:noreply, state}
def handle_cast(:reconnect, %{connect_attempts: connect_attempts} = conn_state)
when connect_attempts > 0 do
{_, conn_state} = Initializer.connect(conn_state)
{:noreply, conn_state}
end
def handle_cast(:reconnect, conn_state) do
{:noreply, conn_state}
end
# When a task is done, it automatically messages the return value to the
# calling process. we can use that opportunity to update the messages in
# flight.
@spec handle_info({reference, {:message_done, NSQ.Message.t(), any}}, state) ::
{:noreply, T.conn_state()}
def handle_info({:message_done, _msg, ret_val}, state) do
state |> MessageHandling.update_conn_stats_on_message_done(ret_val)
{:noreply, state}
end
# ------------------------------------------------------- #
# API Definitions #
# ------------------------------------------------------- #
@spec get_state(pid) :: {:ok, state}
def get_state(pid) when is_pid(pid) do
GenServer.call(pid, :state)
end
@spec get_state(connection) :: {:ok, state}
def get_state({_conn_id, pid} = _connection) do
get_state(pid)
end
@spec close(pid, state) :: any
def close(conn, conn_state \\ nil) do
NSQ.Logger.debug("Closing connection #{inspect(conn)}")
conn_state = conn_state || get_state(conn)
# send a CLS command and expect CLOSE_WAIT in response
{:ok, "CLOSE_WAIT"} = cmd(conn, :cls)
# grace period: poll once per second until zero are in flight
result =
wait_for_zero_in_flight_with_timeout(
conn_state.conn_info_pid,
ConnInfo.conn_id(conn_state),
conn_state.msg_timeout
)
# either way, we're exiting
case result do
:ok ->
NSQ.Logger.warn("#{inspect(conn)}: No more messages in flight. Exiting.")
:timeout ->
NSQ.Logger.error(
"#{inspect(conn)}: Timed out waiting for messages to finish. Exiting anyway."
)
end
Process.exit(self(), :normal)
end
@doc """
Calls the command and waits for a response. If a command shouldn't have a
response, use cmd_noresponse.
"""
@spec cmd(pid, tuple, integer) :: {:ok, binary} | {:error, String.t()}
def cmd(conn_pid, cmd, timeout \\ 5000) do
{:ok, ref} = GenServer.call(conn_pid, {:cmd, cmd, :reply})
receive do
{^ref, data} ->
{:ok, data}
after
timeout ->
{:error, "Command #{cmd} took longer than timeout #{timeout}"}
end
end
@doc """
Calls the command but doesn't expect any response. This is important if the
NSQ command does not in fact generate a response. If you use `cmd` and a
response is sent, it will live forever in the command queue.
"""
@spec cmd_noresponse(pid, tuple) :: {:ok, reference} | {:queued, :nosocket}
def cmd_noresponse(conn, cmd) do
GenServer.call(conn, {:cmd, cmd, :noresponse})
end
# ------------------------------------------------------- #
# Private Functions #
# ------------------------------------------------------- #
@spec wait_for_zero_in_flight(pid, binary) :: any
defp wait_for_zero_in_flight(agent_pid, conn_id) do
[in_flight] = ConnInfo.fetch(agent_pid, conn_id, [:messages_in_flight])
NSQ.Logger.debug("Conn #{inspect(conn_id)}: #{in_flight} still in flight")
if in_flight <= 0 do
:ok
else
:timer.sleep(1000)
wait_for_zero_in_flight(agent_pid, conn_id)
end
end
@spec wait_for_zero_in_flight_with_timeout(pid, binary, integer) :: any
defp wait_for_zero_in_flight_with_timeout(agent_pid, conn_id, timeout) do
try do
Task.async(fn -> wait_for_zero_in_flight(agent_pid, conn_id) end)
|> Task.await(timeout)
catch
:timeout, _ -> :timeout
end
end
end
