eredis

%%
%% eredis_pubsub_client
%%
%% This client implements a subscriber to a Redis pubsub channel. It
%% is implemented in the same way as eredis_client, except channel
%% messages are streamed to the controlling process. Messages are
%% queued and delivered when the client acknowledges receipt.
%%
%% There is one consuming process per eredis_sub_client.
%% @private
-module(eredis_sub_client).
-behaviour(gen_server).
-include("eredis.hrl").
-include("eredis_sub.hrl").
-define(CONNECT_TIMEOUT, 5000).
-define(RECONNECT_SLEEP, 100).

%% API
-export([start_link/1, stop/1]).

%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
         terminate/2, code_change/3]).

%%
%% API
%%

-spec start_link(eredis_sub:sub_options()) ->
          {ok, Pid::pid()} | {error, Reason::term()}.
start_link(Options) ->
    case proplists:lookup(name, Options) of
        {name, Name} ->
            gen_server:start_link(Name, ?MODULE, Options, []);
        none ->
            gen_server:start_link(?MODULE, Options, [])
    end.

stop(Pid) ->
    gen_server:call(Pid, stop).

%%====================================================================
%% gen_server callbacks
%%====================================================================

init(Options) ->
    %% Options
    Host           = proplists:get_value(host, Options, "127.0.0.1"),
    Port           = proplists:get_value(port, Options, 6379),
    Database       = proplists:get_value(database, Options, 0),
    Username       = proplists:get_value(username, Options, undefined),
    Password       = proplists:get_value(password, Options, undefined),
    ReconnectSleep = proplists:get_value(reconnect_sleep, Options, ?RECONNECT_SLEEP),
    ConnectTimeout = proplists:get_value(connect_timeout, Options, ?CONNECT_TIMEOUT),
    SocketOptions  = proplists:get_value(socket_options, Options, []),
    TlsOptions     = proplists:get_value(tls, Options, []),
    Transport      = case TlsOptions of
                         [] -> gen_tcp;
                         _ -> ssl
                     end,
    Active         = eredis_client:get_active(SocketOptions, TlsOptions),

    %% eredis_pub specific options
    MaxQueueSize   = proplists:get_value(max_queue_size, Options, infinity),
    QueueBehaviour = proplists:get_value(queue_behaviour, Options, drop),

    State = #state{host            = Host,
                   port            = Port,
                   database        = eredis_client:read_database(Database),
                   auth_cmd        = eredis_client:get_auth_command(Username,
                                                                    Password),
                   reconnect_sleep = ReconnectSleep,
                   connect_timeout = ConnectTimeout,
                   socket_options  = SocketOptions,
                   tls_options     = TlsOptions,
                   transport       = Transport,
                   active          = Active,
                   channels        = [],
                   pchannels       = [],
                   parser_state    = eredis_parser:init(),
                   msg_queue       = queue:new(),
                   max_queue_size  = MaxQueueSize,
                   queue_behaviour = QueueBehaviour},

    %% Same sync/async connect behaviour as eredis_client.
    case ReconnectSleep of
        no_reconnect ->
            case connect(State) of
                {ok, _NewState} = Res -> Res;
                {error, Reason} -> {stop, Reason}
            end;
        T when is_integer(T) ->
            self() ! initiate_connection,
            {ok, State}
    end.

%% Set the controlling process. All messages on all channels are directed here.
handle_call({controlling_process, Pid}, _From, State) ->
    case State#state.controlling_process of
        undefined ->
            ok;
        {OldRef, _OldPid} ->
            erlang:demonitor(OldRef)
    end,
    Ref = erlang:monitor(process, Pid),
    {reply, ok, State#state{controlling_process={Ref, Pid}, msg_state = ready}};

handle_call(get_channels, _From, State) ->
    {reply, {ok, State#state.channels ++ State#state.pchannels}, State};


handle_call(stop, _From, State) ->
    {stop, normal, ok, State};

handle_call(_Request, _From, State) ->
    {reply, unknown_request, State}.


%% Controlling process acks, but we have no connection. When the
%% connection comes back up, we should be ready to forward a message
%% again.
handle_cast({ack_message, Pid},
            #state{controlling_process={_, Pid}, socket = undefined} = State) ->
    {noreply, State#state{msg_state = ready}};

%% Controlling process acknowledges receipt of previous message. Send
%% the next if there is any messages queued or ask for more on the
%% socket.
handle_cast({ack_message, Pid},
            #state{controlling_process={_, Pid}} = State) ->
    NewState = case queue:out(State#state.msg_queue) of
                   {empty, _Queue} ->
                       State#state{msg_state = ready};
                   {{value, Msg}, Queue} ->
                       send_to_controller(Msg, State),
                       State#state{msg_queue = Queue, msg_state = need_ack}
               end,
    {noreply, NewState};

handle_cast({subscribe, Pid, Channels},
            #state{transport = Transport,
                   socket = Socket,
                   controlling_process = {_, Pid}} = State) ->
    ok = send_subscribe_command(Transport, Socket, "SUBSCRIBE", Channels),
    NewChannels = add_channels(Channels, State#state.channels),
    {noreply, State#state{channels = NewChannels}};


handle_cast({psubscribe, Pid, Channels},
            #state{transport = Transport,
                   socket = Socket,
                   controlling_process = {_, Pid}} = State) ->
    ok = send_subscribe_command(Transport, Socket, "PSUBSCRIBE", Channels),
    NewChannels = add_channels(Channels, State#state.pchannels),
    {noreply, State#state{pchannels = NewChannels}};



handle_cast({unsubscribe, Pid, Channels},
            #state{transport = Transport,
                   controlling_process = {_, Pid}} = State) ->
    Command = eredis:create_multibulk(["UNSUBSCRIBE" | Channels]),
    ok = Transport:send(State#state.socket, Command),
    NewChannels = remove_channels(Channels, State#state.channels),
    {noreply, State#state{channels = NewChannels}};



handle_cast({punsubscribe, Pid, Channels},
            #state{transport = Transport,
                   controlling_process = {_, Pid}} = State) ->
    Command = eredis:create_multibulk(["PUNSUBSCRIBE" | Channels]),
    ok = Transport:send(State#state.socket, Command),
    NewChannels = remove_channels(Channels, State#state.pchannels),
    {noreply, State#state{pchannels = NewChannels}};



handle_cast({ack_message, _}, State) ->
    {noreply, State};

handle_cast(_Msg, State) ->
    {noreply, State}.


%% Receive TCP/TLS data from socket. Match `Socket' to enforce sanity.
handle_info({Type, Socket, Bs}, #state{socket = Socket} = State)
  when Type =:= tcp; Type =:= ssl->
    NewState = handle_response(Bs, State),
    case queue:len(NewState#state.msg_queue) > NewState#state.max_queue_size of
        true ->
            case State#state.queue_behaviour of
                drop ->
                    Msg = {dropped, queue:len(NewState#state.msg_queue)},
                    send_to_controller(Msg, NewState),
                    {noreply, NewState#state{msg_queue = queue:new()}};
                exit ->
                    {stop, max_queue_size, State}
            end;
        false ->
            {noreply, NewState}
    end;

%% Socket switched to passive mode due to {active, N}.
handle_info({Passive, Socket},
            #state{socket = Socket, transport = Transport, active = N} = State)
  when Passive =:= tcp_passive; Passive =:= ssl_passive ->
    case setopts(Socket, Transport, [{active, N}]) of
        ok ->
            {noreply, State};
        {error, Reason} ->
            maybe_reconnect(Reason, State)
    end;

handle_info({Error, Socket, _Reason}, #state{socket = Socket} = State)
  when Error =:= tcp_error; Error =:= ssl_error ->
    maybe_reconnect(Error, State);

%% Socket got closed, for example by Redis terminating idle
%% clients. If desired, spawn of a new process which will try to reconnect and
%% notify us when Redis is ready. In the meantime, we can respond with
%% an error message to all our clients.
handle_info({Closed, Socket}, #state{socket = OurSocket} = State)
  when Closed =:= tcp_closed orelse Closed =:= ssl_closed,
       Socket =:= OurSocket orelse Socket =:= fake_socket ->
    send_to_controller({eredis_disconnected, self()}, State),
    maybe_reconnect(Closed, State#state{socket = undefined});

handle_info(initiate_connection,
            #state{socket = undefined,
                   reconnect_sleep = ReconnectSleep} = State) ->
    case connect(State) of
        {ok, NewState} ->
            {noreply, NewState};
        {error, Reason} ->
            erlang:send_after(ReconnectSleep, self(), {reconnect, Reason}),
            {noreply, State}
    end;

handle_info({reconnect, Reason}, #state{socket = undefined} = State) ->
    %% Scheduled reconnect, if disconnected.
    maybe_reconnect(Reason, State#state{reconnect_timer = undefined});
handle_info({reconnect, _Reason}, State) ->
    %% Already connected.
    {noreply, State#state{reconnect_timer = undefined}};

%% Our controlling process is down.
handle_info({'DOWN', Ref, process, Pid, _Reason},
            #state{controlling_process={Ref, Pid}} = State) ->
    {stop, shutdown, State#state{controlling_process=undefined,
                                 msg_state=ready,
                                 msg_queue=queue:new()}};

%% eredis can be used in Poolboy, but it requires to support a simple API
%% that Poolboy uses to manage the connections.
handle_info(stop, State) ->
    {stop, shutdown, State};

handle_info(_Info, State) ->
    {noreply, State}.

terminate(_Reason, #state{socket = undefined}) ->
    ok;
terminate(_Reason, #state{socket = Socket, transport = Transport}) ->
    Transport:close(Socket).

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%--------------------------------------------------------------------
%%% Internal functions
%%--------------------------------------------------------------------

%% @doc When no channels are given, we unsubscribe from all channels. This
%% matches the semantics of (P)UNSUBSCRIBE without channels.
-spec remove_channels([binary()], [binary()]) -> [binary()].
remove_channels([], _OldChannels) ->
    [];
remove_channels(Channels, OldChannels) ->
    lists:foldl(fun lists:delete/2, OldChannels, Channels).

-spec add_channels([binary()], [binary()]) -> [binary()].
add_channels(Channels, OldChannels) ->
    lists:foldl(fun(C, Cs) ->
                        case lists:member(C, Cs) of
                            true ->
                                Cs;
                            false ->
                                [C|Cs]
                        end
                end, OldChannels, Channels).

%% @doc Sends a subscribe or psubscribe command to Redis.
-spec send_subscribe_command(Transport :: gen_tcp | ssl,
                             Socket :: gen_tcp:socket() | ssl:socket(),
                             Command :: iodata(),
                             Channels :: list()) -> ok.
send_subscribe_command(_Transport, _Socket, _Command, []) ->
    ok;
send_subscribe_command(Transport, Socket, Command, Channels) ->
    Cmd = eredis:create_multibulk([Command | Channels]),
    Transport:send(Socket, Cmd).

-spec handle_response(Data::binary(), State::#state{}) -> NewState::#state{}.
%% @doc: Handle the response coming from Redis. This should only be
%% channel messages that we should forward to the controlling process
%% or queue if the previous message has not been acked. If there are
%% more than a single response in the data we got, queue the responses
%% and serve them up when the controlling process is ready
handle_response(Data, #state{parser_state = ParserState} = State) ->
    case eredis_parser:parse(ParserState, Data) of
        {ReturnCode, Value, NewParserState} ->
            reply({ReturnCode, Value},
                  State#state{parser_state=NewParserState});

        {ReturnCode, Value, Rest, NewParserState} ->
            NewState = reply({ReturnCode, Value},
                             State#state{parser_state=NewParserState}),
            handle_response(Rest, NewState);

        {continue, NewParserState} ->
            State#state{parser_state = NewParserState}
    end.

%% @doc: Sends a reply to the controlling process if the process has
%% acknowledged the previous process, otherwise the message is queued
%% for later delivery.
reply({ok, [<<"message">>, Channel, Message]}, State) ->
    queue_or_send({message, Channel, Message, self()}, State);

reply({ok, [<<"pmessage">>, Pattern, Channel, Message]}, State) ->
    queue_or_send({pmessage, Pattern, Channel, Message, self()}, State);



reply({ok, [<<"subscribe">>, Channel, _]}, State) ->
    queue_or_send({subscribed, Channel, self()}, State);

reply({ok, [<<"psubscribe">>, Channel, _]}, State) ->
    queue_or_send({subscribed, Channel, self()}, State);


reply({ok, [<<"unsubscribe">>, Channel, _]}, State) ->
    queue_or_send({unsubscribed, Channel, self()}, State);


reply({ok, [<<"punsubscribe">>, Channel, _]}, State) ->
    queue_or_send({unsubscribed, Channel, self()}, State);
reply({ReturnCode, Value}, State) ->
    throw({unexpected_response_from_redis, ReturnCode, Value, State}).


queue_or_send(Msg, State) ->
    case State#state.msg_state of
        need_ack ->
            MsgQueue = queue:in(Msg, State#state.msg_queue),
            State#state{msg_queue = MsgQueue};
        ready ->
            send_to_controller(Msg, State),
            State#state{msg_state = need_ack}
    end.


%% @doc: Helper for connecting to Redis. These commands are
%% synchronous and if Redis returns something we don't expect, we
%% crash. Returns {ok, State} or {error, Reason}.
connect(#state{host = Host, port = Port, socket_options = SocketOptions,
               transport = Transport, active = Active,
               connect_timeout = ConnectTimeout, tls_options = TlsOptions,
               auth_cmd = AuthCmd, database = Db} = State) ->
    case eredis_client:connect(Host, Port, SocketOptions, TlsOptions,
                               ConnectTimeout, AuthCmd, Db, Active) of
        {ok, Socket} ->
            %% Re-subscribe to channels. Channels are stored in reverse order in
            %% state.
            ok = send_subscribe_command(Transport, Socket, "SUBSCRIBE",
                                        lists:reverse(State#state.channels)),
            ok = send_subscribe_command(Transport, Socket, "PSUBSCRIBE",
                                        lists:reverse(State#state.pchannels)),

            %% Notify application that connection is ready.
            send_to_controller({eredis_connected, self()}, State),

            %% In case the connection terminates immediately (this happens with
            %% an expired certificate with TLS 1.3) schedule a reconnect already
            %% so that we don't try to reconnect if an error is received before
            %% reconnect_sleep milliseconds has elapsed.
            {ok, schedule_reconnect(unknown, State#state{socket = Socket})};
        Error ->
            Error
    end.

close_socket(#state{socket = undefined} = State) ->
    State;
close_socket(#state{socket = Socket, transport = Transport} = State) ->
    send_to_controller({eredis_disconnected, self()}, State),
    Transport:close(Socket),
    %% Throw away the socket. The absence of a socket is used to
    %% signal we are "down"; discard possibly patrially parsed data
    State#state{socket = undefined, parser_state = eredis_parser:init()}.

%% @doc Schedules a reconnect attempt, if reconnect is enabled.
-spec schedule_reconnect(Reason :: any(), #state{}) -> #state{}.
schedule_reconnect(_Reason, #state{reconnect_sleep = no_reconnect} = State) ->
    State;
schedule_reconnect(Reason, #state{reconnect_sleep = ReconnectSleep,
                                  reconnect_timer = undefined} = State) ->
    TRef = erlang:send_after(ReconnectSleep, self(), {reconnect, Reason}),
    State#state{reconnect_timer = TRef}.

%% @doc Reconnects, but not if a reconnect has been scheduled or if reconnect is
%% disabled. The socket in the state is closed, if any. Returns {noreply, State}
%% or {stop, ExitReason, State} like handle_info.
maybe_reconnect(_Reason, #state{reconnect_sleep = no_reconnect} = State) ->
    %% If we aren't going to reconnect, then there is nothing else for this
    %% process to do.
    {stop, normal, close_socket(State)};
maybe_reconnect(_Reason, #state{reconnect_timer = TRef} = State)
  when is_reference(TRef) ->
    %% Reconnect already scheduled.
    {noreply, close_socket(State)};
maybe_reconnect(_Reason, State) ->
    State1 = close_socket(State),

    %% Controller might want to be notified about every reconnect attempt
    send_to_controller({eredis_reconnect_attempt, self()}, State1),

    case connect(State1) of
        {ok, State2} ->
            {noreply, State2};
        {error, Reason} ->
            %% Controller might want to be notified about every reconnect
            %% failure and reason
            send_to_controller({eredis_reconnect_failed, self(),
                                {error, {connection_error, Reason}}}, State1),
            {noreply, schedule_reconnect(Reason, State1)}
    end.

setopts(Socket, _Transport=gen_tcp, Opts) -> inet:setopts(Socket, Opts);
setopts(Socket, _Transport=ssl, Opts)     ->  ssl:setopts(Socket, Opts).

send_to_controller(_Msg, #state{controlling_process=undefined}) ->
    ok;
send_to_controller(Msg, #state{controlling_process={_Ref, Pid}}) ->
    Pid ! Msg.