%% ========================================================================================================== %% PGPool - A PosgreSQL client that automatically uses connection pools and reconnects in case of errors. %% %% The MIT License (MIT) %% %% Copyright (c) 2016 Roberto Ostinelli and Neato Robotics, Inc. %% %% 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(pgpool_worker). -behaviour(gen_server). -behaviour(poolboy_worker). %% API -export([start_link/1]). -export([squery/2, squery/3]). -export([equery/3, equery/4]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). %% records -record(state, { database_name = undefined :: undefined | atom(), host = undefined :: undefined | string(), user = undefined :: undefined | string(), pass = undefined :: undefined | string(), connect_options = undefined :: undefined | list(), conn = undefined :: undefined | any(), timer_ref = undefined :: undefined | reference() }). %% macros -define(RECONNECT_TIMEOUT_MS, 5000). -define(RETRY_SLEEP_MS, 1000). %% =================================================================== %% API %% =================================================================== -spec start_link(Args :: list()) -> {ok, pid()} | {error, any()}. start_link(Args) -> gen_server:start_link(?MODULE, Args, []). -spec squery(DatabaseName :: atom(), Sql :: string() | iodata()) -> any() | {error, no_connection}. squery(DatabaseName, Sql) -> squery(DatabaseName, Sql, 0). -spec squery(DatabaseName :: atom(), Sql :: string() | iodata(), RetryTimeout :: non_neg_integer() | infinity) -> any() | {error, no_connection}. squery(DatabaseName, Sql, RetryTimeout) -> poolboy:transaction(DatabaseName, fun(Worker) -> case gen_server:call(Worker, {squery, Sql}, infinity) of {error, no_connection} when RetryTimeout =:= infinity -> timer:sleep(?RETRY_SLEEP_MS), squery(DatabaseName, Sql, infinity); {error, no_connection} when RetryTimeout > 0 -> timer:sleep(?RETRY_SLEEP_MS), squery(DatabaseName, Sql, RetryTimeout - ?RETRY_SLEEP_MS); Result -> Result end end). -spec equery(DatabaseName :: atom(), Statement :: string(), Params :: list()) -> any() | {error, no_connection}. equery(DatabaseName, Statement, Params) -> equery(DatabaseName, Statement, Params, 0). -spec equery(DatabaseName :: atom(), Statement :: string(), Params :: list(), RetryTimeout :: non_neg_integer() | infinity) -> any() | {error, no_connection}. equery(DatabaseName, Statement, Params, RetryTimeout) -> poolboy:transaction(DatabaseName, fun(Worker) -> case gen_server:call(Worker, {equery, Statement, Params}, infinity) of {error, no_connection} when RetryTimeout =:= infinity -> timer:sleep(?RETRY_SLEEP_MS), equery(DatabaseName, Statement, Params, infinity); {error, no_connection} when RetryTimeout > 0 -> timer:sleep(?RETRY_SLEEP_MS), equery(DatabaseName, Statement, Params, RetryTimeout - ?RETRY_SLEEP_MS); Result -> Result end end). %% =================================================================== %% Callbacks %% =================================================================== %% ---------------------------------------------------------------------------------------------------------- %% Init %% ---------------------------------------------------------------------------------------------------------- -spec init(Args :: list()) -> {ok, #state{}} | {ok, #state{}, Timeout :: non_neg_integer()} | ignore | {stop, Reason :: any()}. init(Args) -> process_flag(trap_exit, true), %% read options DatabaseName = proplists:get_value(database_name, Args), %% read connection options Host = proplists:get_value(host, Args), User = proplists:get_value(user, Args), Pass = proplists:get_value(pass, Args), ConnectOptions = proplists:get_value(options, Args), %% build state State = #state{ database_name = DatabaseName, host = Host, user = User, pass = Pass, connect_options = ConnectOptions }, %% connect State1 = connect(State), %% return {ok, State1}. %% ---------------------------------------------------------------------------------------------------------- %% Call messages %% ---------------------------------------------------------------------------------------------------------- -spec handle_call(Request :: any(), From :: any(), #state{}) -> {reply, Reply :: any(), #state{}} | {reply, Reply :: any(), #state{}, Timeout :: non_neg_integer()} | {noreply, #state{}} | {noreply, #state{}, Timeout :: non_neg_integer()} | {stop, Reason :: any(), Reply :: any(), #state{}} | {stop, Reason :: any(), #state{}}. handle_call(_Msg, _From, #state{conn = undefined} = State) -> {reply, {error, no_connection}, State}; handle_call({squery, Sql}, _From, #state{conn = Conn} = State) -> {reply, epgsql:squery(Conn, Sql), State}; handle_call({equery, Statement, Params}, _From, #state{conn = Conn} = State) -> {reply, epgsql:equery(Conn, Statement, Params), State}; handle_call(Request, From, State) -> error_logger:warning_msg("Received from ~p an unknown call message: ~p", [Request, From]), {reply, undefined, State}. %% ---------------------------------------------------------------------------------------------------------- %% Cast messages %% ---------------------------------------------------------------------------------------------------------- -spec handle_cast(Msg :: any(), #state{}) -> {noreply, #state{}} | {noreply, #state{}, Timeout :: non_neg_integer()} | {stop, Reason :: any(), #state{}}. handle_cast(Msg, State) -> error_logger:warning_msg("Received an unknown cast message: ~p", [Msg]), {noreply, State}. %% ---------------------------------------------------------------------------------------------------------- %% All non Call / Cast messages %% ---------------------------------------------------------------------------------------------------------- -spec handle_info(Info :: any(), #state{}) -> {noreply, #state{}} | {noreply, #state{}, Timeout :: non_neg_integer()} | {stop, Reason :: any(), #state{}}. handle_info(connect, State) -> State1 = connect(State), {noreply, State1}; handle_info({'EXIT', _From, _Reason}, State) -> error_logger:error_msg("epgsql process died, start a timer to reconnect"), State1 = timeout(State), {noreply, State1}; handle_info(Info, State) -> error_logger:warning_msg("Received an unknown info message: ~p", [Info]), {noreply, State}. %% ---------------------------------------------------------------------------------------------------------- %% Terminate %% ---------------------------------------------------------------------------------------------------------- -spec terminate(Reason :: any(), #state{}) -> terminated. terminate(_Reason, #state{conn = Conn}) -> %% terminate case Conn of undefined -> ok; _ -> ok = epgsql:close(Conn) end, terminated. %% ---------------------------------------------------------------------------------------------------------- %% Convert process state when code is changed. %% ---------------------------------------------------------------------------------------------------------- -spec code_change(OldVsn :: any(), #state{}, Extra :: any()) -> {ok, #state{}}. code_change(_OldVsn, State, _Extra) -> {ok, State}. %% =================================================================== %% Internal %% =================================================================== -spec connect(#state{}) -> #state{}. connect(#state{ database_name = DatabaseName, host = Host, user = User, pass = Pass, connect_options = ConnectOptions } = State) -> error_logger:info_msg("Connecting to database ~p", [DatabaseName]), case epgsql:connect(Host, User, Pass, ConnectOptions) of {ok, Conn} -> State#state{conn = Conn}; Error -> error_logger:error_msg("Error connecting to database ~p with host ~p, user ~p, options ~p: ~p, will try reconnecting in ~p ms", [ DatabaseName, Host, User, ConnectOptions, Error, ?RECONNECT_TIMEOUT_MS ]), State#state{conn = undefined} end. -spec timeout(#state{}) -> #state{}. timeout(#state{ timer_ref = TimerPrevRef } = State) -> case TimerPrevRef of undefined -> ignore; _ -> erlang:cancel_timer(TimerPrevRef) end, TimerRef = erlang:send_after(?RECONNECT_TIMEOUT_MS, self(), connect), State#state{timer_ref = TimerRef}.