%% -*- mode: erlang;erlang-indent-level: 8;indent-tabs-mode:t -*- %% Copyright (c) 2009-2012 %% Bill Warnecke %% Jacob Vorreuter %% Henning Diedrich %% Eonblast Corporation %% %% 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 restric- %% tion, 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(emysql_conn_mgr). -behaviour(gen_server). -export([start_link/0, init/1, handle_call/3, handle_cast/2, handle_info/2]). -export([terminate/2, code_change/3]). -export([pools/0, add_pool/1, has_pool/1,remove_pool/1, add_connections/2, remove_connections/2, lock_connection/1, wait_for_connection/1, wait_for_connection/2, pass_connection/1, replace_connection_as_locked/2, replace_connection_as_available/2, find_pool/2, give_manager_control/1]). -include("emysql.hrl"). -record(state, {pools, lockers = dict:new()}). %%==================================================================== %% API %%==================================================================== %%-------------------------------------------------------------------- %% Function: start_link() -> {ok,Pid} | ignore | {error,Error} %% Description: Starts the server %%-------------------------------------------------------------------- start_link() -> gen_server:start_link({local, ?MODULE}, ?MODULE, [], []). pools() -> gen_server:call(?MODULE, pools, infinity). add_pool(Pool) -> do_gen_call({add_pool, Pool}). has_pool(Pool) -> do_gen_call({has_pool, Pool}). remove_pool(PoolId) -> do_gen_call({remove_pool, PoolId}). add_connections(PoolId, Conns) when is_list(Conns) -> do_gen_call({add_connections, PoolId, Conns}). remove_connections(PoolId, Num) when is_integer(Num) -> do_gen_call({remove_connections, PoolId, Num}). lock_connection(PoolId)-> do_gen_call({lock_connection, PoolId, false, self()}). wait_for_connection(PoolId)-> wait_for_connection(PoolId, lock_timeout()). wait_for_connection(PoolId ,Timeout)-> %% try to lock a connection. if no connections are available then %% wait to be notified of the next available connection %-% io:format("~p waits for connection to pool ~p~n", [self(), PoolId]), case do_gen_call({lock_connection, PoolId, true, self()}) of unavailable -> %-% io:format("~p is queued~n", [self()]), receive {connection, Connection} -> Connection after Timeout -> do_gen_call({abort_wait, PoolId}), receive {connection, Connection} -> Connection after 0 -> exit(connection_lock_timeout) end end; Connection -> %-% io:format("~p gets connection~n", [self()]), Connection end. pass_connection(Connection) -> do_gen_call({{replace_connection, available}, Connection, Connection}). replace_connection_as_available(OldConn, NewConn) -> do_gen_call({{replace_connection, available}, OldConn, NewConn}). replace_connection_as_locked(OldConn, NewConn) -> do_gen_call({{replace_connection, locked}, OldConn, NewConn}). give_manager_control(Socket) -> case whereis(?MODULE) of undefined -> {error ,failed_to_find_conn_mgr}; MgrPid -> gen_tcp:controlling_process(Socket, MgrPid) end. %% the stateful loop functions of the gen_server never %% want to call exit/1 because it would crash the gen_server. %% instead we want to return error tuples and then throw %% the error once outside of the gen_server process do_gen_call(Msg) -> case gen_server:call(?MODULE, Msg, infinity) of {error, Reason} -> exit(Reason); Result -> Result end. %%==================================================================== %% gen_server callbacks %%==================================================================== %%-------------------------------------------------------------------- %% Function: init(Args) -> {ok, State} | %% {ok, State, Timeout} | %% ignore | %% {stop, Reason} %% Description: Initiates the server %%-------------------------------------------------------------------- init([]) -> Pools = lists:map( fun (Pool1) -> case emysql_conn:open_connections(Pool1) of {ok, Pool2} -> Pool2; {error, Reason} -> throw(Reason) end end, initialize_pools() ), {ok, #state{pools=Pools}}. %%-------------------------------------------------------------------- %% 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(pools, _From, State) -> {reply, State#state.pools, State}; handle_call({has_pool, PoolID}, _From, State) -> case find_pool(PoolID, State#state.pools) of {_, _} -> {reply, true, State}; undefined -> {reply, false, State} end; handle_call({add_pool, Pool}, _From, State) -> case find_pool(Pool#pool.pool_id, State#state.pools) of {_, _} -> {reply, {error, pool_already_exists}, State}; undefined -> {reply, ok, State#state{pools = [Pool|State#state.pools]}} end; handle_call({remove_pool, PoolId}, _From, State) -> case find_pool(PoolId, State#state.pools) of {Pool, OtherPools} -> {reply, Pool, State#state{pools=OtherPools}}; undefined -> {reply, {error, pool_not_found}, State} end; handle_call({add_connections, PoolId, Conns}, _From, State) -> case find_pool(PoolId, State#state.pools) of {Pool, OtherPools} -> Pool1 = Pool#pool{available = queue:join(queue:from_list(Conns), Pool#pool.available)}, {NewPool, NewMonitorData} = serve_waiting_pids(Pool1), NewLockers = lists:foldl(fun({MonitorRef, ConnInfo}, Dict) -> dict:store(MonitorRef, ConnInfo, Dict) end, State#state.lockers, NewMonitorData), State1 = State#state{pools = [NewPool|OtherPools], lockers = NewLockers}, {reply, ok, State1}; undefined -> {reply, {error, pool_not_found}, State} end; handle_call({remove_connections, PoolId, Num}, _From, State) -> case find_pool(PoolId, State#state.pools) of {Pool, OtherPools} -> case Num > queue:len(Pool#pool.available) of true -> State1 = State#state{pools = [Pool#pool{available = queue:new()}]}, {reply, queue:to_list(Pool#pool.available), State1}; false -> {Conns, OtherConns} = queue:split(Num, Pool#pool.available), State1 = State#state{pools = [Pool#pool{available = OtherConns}|OtherPools]}, {reply, queue:to_list(Conns), State1} end; undefined -> {reply, {error, pool_not_found}, State} end; handle_call({lock_connection, PoolId, Wait, Who}, {From, _Mref}, State) -> case find_pool(PoolId, State#state.pools) of {Pool, OtherPools} -> case lock_next_connection(Pool, Who) of {ok, Connection, PoolNow, {MonitorRef, Data}} -> Lockers = State#state.lockers, {reply, Connection, State#state{pools=[PoolNow|OtherPools], lockers = dict:store(MonitorRef, Data, Lockers)}}; unavailable when Wait =:= true -> %% place the calling pid at the end of the waiting queue of its pool PoolNow = Pool#pool{waiting = queue:in(From, Pool#pool.waiting)}, {reply, unavailable, State#state{pools=[PoolNow|OtherPools]}}; unavailable when Wait =:= false -> {reply, unavailable, State} end; undefined -> {reply, {error, pool_not_found}, State} end; handle_call({abort_wait, PoolId}, {From, _Mref}, State) -> case find_pool(PoolId, State#state.pools) of {Pool, OtherPools} -> %% Remove From from the wait queue QueueNow = queue:filter( fun(Pid) -> Pid =/= From end, Pool#pool.waiting), PoolNow = Pool#pool{ waiting = QueueNow }, %% See if the length changed to know if From was removed. OldLen = queue:len(Pool#pool.waiting), NewLen = queue:len(QueueNow), if OldLen =:= NewLen -> Reply = not_waiting; true -> Reply = ok end, {reply, Reply, State#state{pools=[PoolNow|OtherPools]}}; undefined -> {reply, {error, pool_not_found}, State} end; handle_call({{replace_connection, Kind}, OldConn, NewConn}, _From, State) -> %% if an error occurs while doing work over a connection then %% the connection must be closed and a new one created in its %% place. The calling process is responsible for creating the %% new connection, closing the old one and replacing it in state. %% This function expects a new, available connection to be %% passed in to serve as the replacement for the old one. %% But i.e. if the sql server is down, it can be fed a dead %% old connection as new connection, to preserve the pool size. case find_pool(OldConn#emysql_connection.pool_id, State#state.pools) of {#pool{available = Available, locked = Locked} = Pool, Pools} -> OldRef = OldConn#emysql_connection.monitor_ref, Stripped = gb_trees:delete_any(OldConn#emysql_connection.id, Locked), {NewPool, NewMonitors} = case Kind of available -> serve_waiting_pids( Pool#pool { locked = Stripped, available = queue:in( NewConn#emysql_connection { locked_at = undefined, monitor_ref = undefined}, Available) }); locked -> {Pool#pool { locked = gb_trees:enter(NewConn#emysql_connection.id, NewConn#emysql_connection{monitor_ref = OldRef}, Stripped) }, []} % There are no new monitors set up end, Lockers = State#state.lockers, NewLockers = case Kind of available -> %% Remove the current monitor as the connection has been freed %% Add the monitors for pids that has been take from the waiting queue erlang:demonitor(OldRef, [flush]), lists:foldl(fun({MonitorRef, ConnInfo}, Dict) -> dict:store(MonitorRef, ConnInfo, Dict) end, dict:erase(OldRef, Lockers), NewMonitors); locked -> %% We need to keep the monitor here dict:store(OldRef, {NewConn#emysql_connection.pool_id, NewConn#emysql_connection.id}, Lockers) end, {reply, ok, State#state{pools = [NewPool|Pools], lockers = NewLockers}}; undefined -> {reply, {error, pool_not_found}, State} end; handle_call(_, _From, State) -> {reply, {error, invalid_call}, State}. %%-------------------------------------------------------------------- %% Function: handle_cast(Msg, State) -> {noreply, State} | %% {noreply, State, Timeout} | %% {stop, Reason, State} %% Description: Handling cast messages %%-------------------------------------------------------------------- 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({'DOWN', MonitorRef, _, _, _}, State) -> case dict:find(MonitorRef, State#state.lockers) of {ok, {PoolId, ConnId}} -> case find_pool(PoolId, State#state.pools) of {Pool, _} -> case gb_trees:lookup(ConnId, Pool#pool.locked) of {value, Conn} -> async_reset_conn(State#state.pools, Conn); _ -> ok end; _ -> ok end; _ -> ok end, {noreply, State}; handle_info(_Info, State) -> {noreply, State}. async_reset_conn(Pools, Conn) -> spawn(fun() -> %% This interacts with the conn mgr so needs to be spawned %% TODO: refactor emysql_conn:reset_connection(Pools, Conn, pass) end). %%-------------------------------------------------------------------- %% 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(_OldVsn, State, _Extra) -> {ok, State}. %%-------------------------------------------------------------------- %%% Internal functions %%-------------------------------------------------------------------- initialize_pools() -> %% if the emysql application values are not present in the config %% file we will initialize and empty set of pools. Otherwise, the %% values defined in the config are used to initialize the state. [ #pool{ pool_id = PoolId, size = proplists:get_value(size, Props, 1), user = proplists:get_value(user, Props), password = proplists:get_value(password, Props), host = proplists:get_value(host, Props), port = proplists:get_value(port, Props), database = proplists:get_value(database, Props), encoding = proplists:get_value(encoding, Props), start_cmds = proplists:get_value(start_cmds, Props, []) } || {PoolId, Props} <- emysql_app:pools() ]. find_pool(PoolId, Pools) -> find_pool(PoolId, Pools, []). find_pool(_, [], _) -> undefined; find_pool(PoolId, [#pool{pool_id = PoolId} = Pool|Tail], OtherPools) -> {Pool, lists:append(OtherPools, Tail)}; find_pool(PoolId, [Pool|Tail], OtherPools) -> find_pool(PoolId, Tail, [Pool|OtherPools]). lock_next_connection(Pool, Who) -> case lock_next_connection(Pool#pool.available, Pool#pool.locked, Who) of {ok, Connection, OtherAvailable, NewLocked, MonitorTuple} -> {ok ,Connection ,Pool#pool{available=OtherAvailable, locked=NewLocked}, MonitorTuple}; unavailable -> unavailable end. lock_next_connection(Available ,Locked, Who) -> case queue:out(Available) of {{value, Conn}, OtherAvailable} -> MonitorRef = erlang:monitor(process, Who), NewConn = connection_locked_at(Conn, MonitorRef), MonitorTuple = {MonitorRef, {NewConn#emysql_connection.pool_id, NewConn#emysql_connection.id}}, NewLocked = gb_trees:enter(NewConn#emysql_connection.id, NewConn, Locked), {ok, NewConn, OtherAvailable, NewLocked, MonitorTuple}; {empty, _} -> unavailable end. connection_locked_at(Conn, MonitorRef) -> Conn#emysql_connection{locked_at=lists:nth(2, tuple_to_list(now())), monitor_ref = MonitorRef}. serve_waiting_pids(Pool) -> {Waiting, Available, Locked, NewRefs} = serve_waiting_pids(Pool#pool.waiting, Pool#pool.available, Pool#pool.locked, []), {Pool#pool{waiting=Waiting, available=Available, locked=Locked}, NewRefs}. serve_waiting_pids(Waiting, Available, Locked, MonitorRefs) -> case queue:is_empty(Waiting) of false -> Who = queue:get(Waiting), case lock_next_connection(Available, Locked, Who) of {ok, Connection, OtherAvailable, NewLocked, NewRef} -> {{value, Pid}, OtherWaiting} = queue:out(Waiting), case erlang:is_process_alive(Pid) of true -> erlang:send(Pid, {connection, Connection}), serve_waiting_pids(OtherWaiting, OtherAvailable, NewLocked, [NewRef | MonitorRefs]); _ -> serve_waiting_pids(OtherWaiting, Available, Locked, MonitorRefs) end; unavailable -> {Waiting, Available, Locked, MonitorRefs} end; true -> {Waiting, Available, Locked, MonitorRefs} end. lock_timeout() -> emysql_app:lock_timeout().