%% Copyright (c) 2011 Jacob Vorreuter %% %% Permission is hereby granted, free of charge, to any person %% obtaining a copy of this software and associated documentation %% files (the "Software"), to deal in the Software without %% restriction, including without limitation the rights to use, %% copy, modify, merge, publish, distribute, sublicense, and/or sell %% copies of the Software, and to permit persons to whom the %% Software is furnished to do so, subject to the following %% conditions: %% %% The above copyright notice and this permission notice shall be %% included in all copies or substantial portions of the Software. %% %% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, %% EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES %% OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND %% NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT %% HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, %% WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING %% FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR %% OTHER DEALINGS IN THE SOFTWARE. -module(redo). -behaviour(gen_server). -include("redo_logging.hrl"). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). -export([start_link/0, start_link/1, start_link/2, cmd/1, cmd/2, cmd/3, subscribe/1, subscribe/2, shutdown/1, async_cmd/1, async_cmd/2]). -record(state, {host, port, pass, db, sock, queue, subscriber, cancelled, acc, buffer, reconnect}). -define(TIMEOUT, 30000). -spec start_link() -> {ok, pid()} | {error, term()}. start_link() -> start_link([]). -spec start_link(atom() | list()) -> {ok, pid()} | {error, term()}. start_link(Name) when is_atom(Name) -> start_link(Name, []); start_link(Opts) when is_list(Opts) -> start_link(?MODULE, Opts). -spec start_link(atom(), list()) -> {ok, pid()} | {error, term()}. start_link(undefined, Opts) when is_list(Opts) -> gen_server:start_link(?MODULE, [Opts], []); start_link(Name, Opts) when is_atom(Name), is_list(Opts) -> gen_server:start_link({local, Name}, ?MODULE, [Opts], []). cmd(Cmd) -> cmd(?MODULE, Cmd, ?TIMEOUT). cmd(NameOrPid, Cmd) -> cmd(NameOrPid, Cmd, ?TIMEOUT). -type redis_value() :: integer() | binary(). -type response() :: redis_value() | [redis_value()] | {'error', Reason::term()}. -spec cmd(atom() | pid(), list() | binary(), integer()) -> response() | [response()]. cmd(NameOrPid, Cmd, Timeout) when is_integer(Timeout); Timeout =:= infinity -> %% format commands to be sent to redis Packets = redo_redis_proto:package(Cmd), %% send the commands and receive back %% unique refs for each packet sent Refs = gen_server:call(NameOrPid, {cmd, Packets}, 2000), receive_resps(NameOrPid, Refs, Timeout). async_cmd(Cmd) -> async_cmd(?MODULE, Cmd). async_cmd(NameOrPid, Cmd) -> Packets = redo_redis_proto:package(Cmd), gen_server:call(NameOrPid, {cmd, Packets}). receive_resps(_NameOrPid, {error, Err}, _Timeout) -> {error, Err}; receive_resps(NameOrPid, [Ref], Timeout) -> %% for a single packet, receive a single reply receive_resp(NameOrPid, Ref, Timeout); receive_resps(NameOrPid, Refs, Timeout) -> %% for multiple packets, build a list of replies [receive_resp(NameOrPid, Ref, Timeout) || Ref <- Refs]. receive_resp(NameOrPid, Ref, Timeout) -> receive %% the connection to the redis server was closed {Ref, closed} -> {error, closed}; {Ref, Val} -> Val %% after the timeout expires, cancel the command and return after Timeout -> gen_server:cast(NameOrPid, {cancel, Ref}), {error, timeout} end. -spec subscribe(list() | binary()) -> reference() | {error, term()}. subscribe(Channel) -> subscribe(?MODULE, Channel). -spec subscribe(atom() | pid(), list() | binary()) -> reference() | {error, term()}. subscribe(NameOrPid, Channel) -> Packet = redo_redis_proto:package(["SUBSCRIBE", Channel]), gen_server:call(NameOrPid, {subscribe, Packet}, 2000). shutdown(NameOrPid) -> gen_server:cast(NameOrPid, shutdown). %%==================================================================== %% gen_server callbacks %%==================================================================== %%-------------------------------------------------------------------- %% Function: init(Args) -> {ok, State} | %% {ok, State, Timeout} | %% ignore | %% {stop, Reason} %% Description: Initiates the server %%-------------------------------------------------------------------- init([Opts]) -> State = init_state(Opts), case connect(State) of State1 when is_record(State1, state) -> {ok, State1}; Err -> {stop, Err} end. %%-------------------------------------------------------------------- %% Function: %% handle_call(Request, From, State) -> {reply, Reply, State} | %% {reply, Reply, State, Timeout} | %% {noreply, State} | %% {noreply, State, Timeout} | %% {stop, Reason, Reply, State} | %% {stop, Reason, State} %% Description: Handling call messages %%-------------------------------------------------------------------- handle_call({cmd, Packets}, {From, _Ref}, #state{subscriber=undefined, queue=Queue}=State) -> case test_connection(State) of State1 when is_record(State1, state) -> %% send each packet to redis %% and generate a unique ref per packet Refs = lists:foldl( fun(Packet, Refs) when is_list(Refs) -> case gen_tcp:send(State1#state.sock, Packet) of ok -> [erlang:make_ref()|Refs]; Err -> Err end; (_Packet, Err) -> Err end, [], Packets), %% enqueue the client pid/refs case Refs of List when is_list(List) -> Refs1 = lists:reverse(Refs), Queue1 = lists:foldl( fun(Ref, Acc) -> queue:in({From, Ref}, Acc) end, Queue, Refs1), {reply, Refs1, State1#state{queue=Queue1}}; Err -> {stop, Err, State1} end; _Err -> %% failed to connect, retry {reply, {error, closed}, State#state{sock=undefined}, 1000} end; handle_call({cmd, _Packets}, _From, State = #state{}) -> {reply, {error, subscriber_mode}, State}; handle_call({subscribe, Packet}, {From, _Ref}, State = #state{}) -> case test_connection(State) of State1 when is_record(State1, state) -> case gen_tcp:send(State1#state.sock, Packet) of ok -> Ref = erlang:make_ref(), {reply, Ref, State1#state{subscriber={From, Ref}}}; Err -> {stop, Err, State1} end; _Err -> %% failed to connect, retry {reply, {error, closed}, State#state{sock=undefined}, 1000} end; %% state doesn't match. Likely outdated or an error. %% Moving from 1.1.0 to here sees the addition of one field handle_call(Msg, From, OldState) when 1+tuple_size(OldState) =:= tuple_size(#state{}), state =:= element(1,OldState) -> {ok, State} = code_change("v1.1.0", OldState, internal_detection), handle_call(Msg, From, State); handle_call(_Msg, _From, State = #state{}) -> {reply, unknown_message, State}. %%-------------------------------------------------------------------- %% Function: handle_cast(Msg, State) -> {noreply, State} | %% {noreply, State, Timeout} | %% {stop, Reason, State} %% Description: Handling cast messages %%-------------------------------------------------------------------- handle_cast(shutdown, State = #state{sock=Sock}) -> case is_port(Sock) of true -> catch gen_tcp:close(Sock); false -> ok end, State1 = close_connection(State#state{sock=undefined}), {stop, shutdown, State1}; handle_cast({cancel, Ref}, #state{cancelled=Cancelled}=State) -> {noreply, State#state{cancelled=[Ref|Cancelled]}}; %% state doesn't match. Likely outdated or an error. %% Moving from 1.1.0 to here sees the addition of one field handle_cast(Msg, OldState) when 1+tuple_size(OldState) =:= tuple_size(#state{}), state =:= element(1,OldState) -> {ok, State} = code_change("v1.1.0", OldState, internal_detection), handle_cast(Msg, State); handle_cast(_Msg, State) -> {noreply, State}. %%-------------------------------------------------------------------- %% Function: handle_info(Info, State) -> {noreply, State} | %% {noreply, State, Timeout} | %% {stop, Reason, State} %% Description: Handling all non call/cast messages %%-------------------------------------------------------------------- handle_info({tcp, Sock, Data}, #state{sock=Sock, buffer=Buffer}=State) -> %% compose the packet to be processed by combining %% the leftover buffer with the new data packet Packet = packet(Buffer, Data), case process_packet(State, Packet) of {ok, State1} -> %% accept next incoming packet inet:setopts(Sock, [{active, once}]), {noreply, State1}; Err -> {stop, Err, State} end; handle_info({tcp_closed, Sock}, #state{sock=Sock, reconnect=false}=State) -> State1 = close_connection(State), {stop, tcp_closed, State1}; handle_info({tcp_closed, Sock}, #state{sock=Sock}=State) -> State1 = close_connection(State), %% reconnect to redis case connect(State1) of State2 when is_record(State2, state) -> {noreply, State2}; _Err -> {noreply, State1#state{sock=undefined}, 1000} end; handle_info({tcp_error, Sock, Reason}, #state{sock=Sock}=State) -> ?WARN("at=tcp_error sock=~p reason=~p", [Sock, Reason]), CState = close_connection(State#state{sock=undefined}), {stop, normal, CState}; %% attempt to reconnect to redis handle_info(timeout, State = #state{}) -> case connect(State) of State1 when is_record(State1, state) -> {noreply, State1}; _Err -> {noreply, State#state{sock=undefined}, 1000} end; %% state doesn't match. Likely outdated or an error. %% Moving from 1.1.0 to here sees the addition of one field handle_info(Msg, OldState) when 1+tuple_size(OldState) =:= tuple_size(#state{}), state =:= element(1,OldState) -> {ok, State} = code_change("v1.1.0", OldState, internal_detection), handle_info(Msg, State); handle_info(_Info, State) -> {noreply, State}. %%-------------------------------------------------------------------- %% Function: terminate(Reason, State) -> void() %% Description: This function is called by a gen_server when it is about to %% terminate. It should be the opposite of Module:init/1 and do any necessary %% cleaning up. When it returns, the gen_server terminates with Reason. %% The return value is ignored. %%-------------------------------------------------------------------- terminate(_Reason, _State) -> ok. %%-------------------------------------------------------------------- %% Func: code_change(OldVsn, State, Extra) -> {ok, NewState} %% Description: Convert process state when code is changed %%-------------------------------------------------------------------- code_change("v1.1.0", OldState, internal_detection) -> {state, Host, Port, Pass, Db, Sock, Queue, Subscriber, Cancelled, Acc, Buffer} = OldState, %% we make the reconnection set to 'true' given that's the default %% value set in init_state/1 {ok, #state{host=Host, port=Port, pass=Pass, db=Db, sock=Sock, queue=Queue, subscriber=Subscriber, cancelled=Cancelled, acc=Acc, buffer=Buffer, reconnect=true}}; code_change(_OldVsn, State, _Extra) -> {ok, State}. %%-------------------------------------------------------------------- %% Internal functions %%-------------------------------------------------------------------- init_state(Opts) -> Host = proplists:get_value(host, Opts, "localhost"), Port = proplists:get_value(port, Opts, 6379), Pass = proplists:get_value(pass, Opts), Db = proplists:get_value(db, Opts, 0), Recn = proplists:get_value(reconnect, Opts, true), #state{ host = Host, port = Port, pass = Pass, db = Db, queue = queue:new(), cancelled = [], acc = [], buffer = {raw, <<>>}, reconnect = Recn }. connect(#state{host=Host, port=Port, pass=Pass, db=Db}=State) -> case connect_socket(Host, Port) of {ok, Sock} -> case auth(Sock, Pass) of ok -> case select_db(Sock, Db) of ok -> inet:setopts(Sock, [{active, once}]), State#state{sock=Sock}; Err -> Err end; Err -> Err end; Err -> Err end. connect_socket(Host, Port) -> SockOpts = [binary, {active, false}, {keepalive, true}, {nodelay, true}], gen_tcp:connect(Host, Port, SockOpts). auth(_Sock, Pass) when Pass == <<>>; Pass == undefined -> ok; auth(Sock, Pass) -> case gen_tcp:send(Sock, [<<"AUTH ">>, Pass, <<"\r\n">>]) of ok -> case gen_tcp:recv(Sock, 0) of {ok, <<"+OK\r\n">>} -> ok; {ok, Err} -> {error, Err}; Err -> Err end; Err -> Err end. select_db(_Sock, 0) -> ok; select_db(Sock, Db) -> case gen_tcp:send(Sock, hd(redo_redis_proto:package(["SELECT", Db]))) of ok -> case gen_tcp:recv(Sock, 0) of {ok, <<"+OK\r\n">>} -> ok; {ok, Err} -> {error, Err}; Err -> Err end; Err -> Err end. test_connection(#state{sock=undefined}=State) -> connect(State); test_connection(State) -> State. process_packet(#state{acc=Acc, queue=Queue, subscriber=Subscriber}=State, Packet) -> case redo_redis_proto:parse(Acc, Packet) of %% the reply has been received in its entirety {ok, Result, Rest} -> case queue:out(Queue) of {{value, {Pid, Ref}}, Queue1} -> send_response(Pid, Ref, Result, Rest, State, Queue1); {empty, Queue1} -> case Subscriber of {Pid, Ref} -> send_response(Pid, Ref, Result, Rest, State, Queue1); undefined -> {error, queue_empty} end end; %% the current reply packet ended abruptly %% we must wait for the next data packet {eof, Acc1, Rest} -> {ok, State#state{acc=Acc1, buffer=Rest}} end. send_response(Pid, Ref, Result, Rest, State, Queue) -> Pid ! {Ref, Result}, case Rest of {raw, <<>>} -> %% we have reached the end of this tcp packet %% wait for the next incoming packet {ok, State#state{acc=[], queue=Queue, buffer = {raw, <<>>}}}; _ -> %% there is still data left in this packet %% we may begin processing the next reply process_packet(State#state{acc=[], queue=Queue}, packet(Rest, <<>>)) end. packet({raw, Buffer}, Data) -> {raw, <>}; packet({multi_bulk, N, Buffer}, Data) -> {multi_bulk, N, <>}. close_connection(State = #state{queue=Queue}) -> %% notify all waiting pids that the connection is closed %% so that they may try resending their requests [Pid ! {Ref, closed} || {Pid, Ref} <- queue:to_list(Queue)], %% notify subscriber pid of disconnect case State#state.subscriber of {Pid, Ref} -> Pid ! {Ref, closed}; _ -> ok end, %% reset the state State#state{ queue = queue:new(), cancelled = [], buffer = {raw, <<>>} }.