From e131db751368179a138315c48720b9a585fab985 Mon Sep 17 00:00:00 2001 From: Thierry Bomandouki Date: Thu, 23 May 2013 12:42:53 +0200 Subject: [PATCH] Allowing eredis to connect to redis servers by Nutcracker proxy (ex Twemproxy) --- src/eredis.erl | 24 ++++++++++++++++++++++-- src/eredis_client.erl | 32 +++++++++++++++++++------------- 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/src/eredis.erl b/src/eredis.erl index 59b42f49..6962c570 100644 --- a/src/eredis.erl +++ b/src/eredis.erl @@ -16,7 +16,7 @@ -define(TIMEOUT, 5000). -export([start_link/0, start_link/1, start_link/2, start_link/3, start_link/4, - start_link/5, stop/1, q/2, q/3, qp/2, qp/3]). + start_link/5, start_link/6, stop/1, q/2, q/3, qp/2, qp/3]). %% Exported for testing -export([create_multibulk/1]). @@ -31,12 +31,18 @@ start_link() -> start_link(Host, Port) -> start_link(Host, Port, 0, ""). +start_link(Host, Port, no_dbselection) -> + start_link(Host, Port, 0, "", no_dbselection); start_link(Host, Port, Database) -> start_link(Host, Port, Database, ""). +start_link(Host, Port, Database, no_dbselection) -> + start_link(Host, Port, Database, "", no_dbselection); start_link(Host, Port, Database, Password) -> start_link(Host, Port, Database, Password, 100). +start_link(Host, Port, Database, Password, no_dbselection) -> + start_link(Host, Port, Database, Password, 100, no_dbselection); start_link(Host, Port, Database, Password, ReconnectSleep) when is_list(Host); is_integer(Port); @@ -46,16 +52,30 @@ start_link(Host, Port, Database, Password, ReconnectSleep) eredis_client:start_link(Host, Port, Database, Password, ReconnectSleep). +start_link(Host, Port, Database, Password, ReconnectSleep, no_dbselection) + when is_list(Host); + is_integer(Port); + is_integer(Database); + is_list(Password); + is_integer(ReconnectSleep) orelse ReconnectSleep =:= no_reconnect -> + + eredis_client:start_link(Host, Port, Database, Password, ReconnectSleep, no_dbselection). %% @doc: Callback for starting from poolboy -spec start_link(server_args()) -> {ok, Pid::pid()} | {error, Reason::term()}. +start_link(no_dbselection) -> + start_link("127.0.0.1", 6379, 0, "", no_dbselection); start_link(Args) -> Host = proplists:get_value(host, Args, "127.0.0.1"), Port = proplists:get_value(port, Args, 6379), Database = proplists:get_value(database, Args, 0), Password = proplists:get_value(password, Args, ""), ReconnectSleep = proplists:get_value(reconnect_sleep, Args, 100), - start_link(Host, Port, Database, Password, ReconnectSleep). + case proplists:is_defined(reconnect_sleep, Args) of + true -> start_link(Host, Port, Database, Password, ReconnectSleep, no_dbselection); + _ -> start_link(Host, Port, Database, Password, ReconnectSleep) + end. + stop(Client) -> eredis_client:stop(Client). diff --git a/src/eredis_client.erl b/src/eredis_client.erl index 3d17be4d..385113c7 100644 --- a/src/eredis_client.erl +++ b/src/eredis_client.erl @@ -28,7 +28,7 @@ -include("eredis.hrl"). %% API --export([start_link/5, stop/1, select_database/2]). +-export([start_link/5, start_link/6, stop/1, select_database/2]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, @@ -43,7 +43,8 @@ socket :: port() | undefined, parser_state :: #pstate{} | undefined, - queue :: queue() | undefined + queue :: queue() | undefined, + no_dbselection :: atom() | undefined }). %% @@ -57,8 +58,10 @@ ReconnectSleep::reconnect_sleep()) -> {ok, Pid::pid()} | {error, Reason::term()}. start_link(Host, Port, Database, Password, ReconnectSleep) -> - gen_server:start_link(?MODULE, [Host, Port, Database, Password, ReconnectSleep], []). + gen_server:start_link(?MODULE, [Host, Port, Database, Password, ReconnectSleep, none], []). +start_link(Host, Port, Database, Password, ReconnectSleep, no_dbselection) -> + gen_server:start_link(?MODULE, [Host, Port, Database, Password, ReconnectSleep, no_dbselection], []). stop(Pid) -> gen_server:call(Pid, stop). @@ -67,7 +70,7 @@ stop(Pid) -> %% gen_server callbacks %%==================================================================== -init([Host, Port, Database, Password, ReconnectSleep]) -> +init([Host, Port, Database, Password, ReconnectSleep, DbSelection]) -> State = #state{host = Host, port = Port, database = list_to_binary(integer_to_list(Database)), @@ -76,8 +79,7 @@ init([Host, Port, Database, Password, ReconnectSleep]) -> parser_state = eredis_parser:init(), queue = queue:new()}, - - case connect(State) of + case connect(State, DbSelection) of {ok, NewState} -> {ok, NewState}; {error, Reason} -> @@ -234,16 +236,20 @@ reply(Value, Queue) -> %% the correct database. These commands are synchronous and if Redis %% returns something we don't expect, we crash. Returns {ok, State} or %% {SomeError, Reason}. -connect(State) -> +connect(State, DbSelection) -> case gen_tcp:connect(State#state.host, State#state.port, ?SOCKET_OPTS) of {ok, Socket} -> case authenticate(Socket, State#state.password) of ok -> - case select_database(Socket, State#state.database) of - ok -> - {ok, State#state{socket = Socket}}; - {error, Reason} -> - {error, {select_error, Reason}} + case DbSelection of + no_dbselection -> {ok, State#state{socket = Socket, no_dbselection = DbSelection}}; + _ -> + case select_database(Socket, State#state.database) of + ok -> + {ok, State#state{socket = Socket, no_dbselection = DbSelection}}; + {error, Reason} -> + {error, {select_error, Reason}} + end end; {error, Reason} -> {error, {authentication_error, Reason}} @@ -282,7 +288,7 @@ do_sync_command(Socket, Command) -> %% successfully issuing the auth and select calls. When we have a %% connection, give the socket to the redis client. reconnect_loop(Client, #state{reconnect_sleep = ReconnectSleep} = State) -> - case catch(connect(State)) of + case catch(connect(State, State#state.no_dbselection)) of {ok, #state{socket = Socket}} -> gen_tcp:controlling_process(Socket, Client), Client ! {connection_ready, Socket};