%%% -*- erlang -*- %%% %%% This file is part of hackney released under the Apache 2 license. %%% See the NOTICE for more information. %%% %%% Copyright (c) 2009, Erlang Training and Consulting Ltd. %%% Copyright (c) 2012-2015, Benoît Chesneau %% @doc pool of sockets connections %% -module(hackney_pool). -behaviour(gen_server). %% PUBLIC API -export([start/0, checkout/4, checkin/2]). -export([ get_stats/1, start_pool/2, stop_pool/1, find_pool/1, notify/2 ]). -export([count/1, count/2, max_connections/1, set_max_connections/2, timeout/1, set_timeout/2, child_spec/2]). -export([start_link/2]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, code_change/3, terminate/2]). -include("hackney.hrl"). -include_lib("hackney_internal.hrl"). -record(state, { name, metrics, max_connections, timeout, clients = dict:new(), queues = dict:new(), % Dest => queue of Froms, pending = dict:new(), connections = dict:new(), sockets = dict:new(), nb_waiters=0}). start() -> %% NB this is first called from hackney_sup:start_link %% BEFORE the hackney_pool ETS table exists ok. %% @doc fetch a socket from the pool checkout(Host0, Port, Transport, #client{options=Opts}=Client) -> Host = string:to_lower(Host0), Pid = self(), RequestRef = Client#client.request_ref, Name = proplists:get_value(pool, Opts, default), Pool = find_pool(Name, Opts), ConnectTimeout = proplists:get_value(connect_timeout, Opts, 8000), case catch gen_server:call(Pool, {checkout, {Host, Port, Transport}, Pid, RequestRef}, ConnectTimeout) of {ok, Socket, Owner} -> CheckinReference = {Host, Port, Transport}, {ok, {Name, RequestRef, CheckinReference, Owner, Transport}, Socket}; {error, no_socket, Owner} -> CheckinReference = {Host, Port, Transport}, {error, no_socket, {Name, RequestRef, CheckinReference, Owner, Transport}}; {error, Reason} -> {error, Reason}; {'EXIT', {timeout, _}} -> % socket will still checkout so to avoid deadlock we send in a cancellation gen_server:cast(Pool, {checkout_cancel, {Host, Port, Transport}, RequestRef}), {error, connect_timeout} end. %% @doc release a socket in the pool checkin({_Name, Ref, Dest, Owner, Transport}, Socket) -> Transport:setopts(Socket, [{active, false}]), case sync_socket(Transport, Socket) of true -> case Transport:controlling_process(Socket, Owner) of ok -> gen_server:call(Owner, {checkin, Ref, Dest, Socket, Transport}, infinity); _Error -> catch Transport:close(Socket), ok end; false -> catch Transport:close(Socket), ok end. get_stats(Pool) -> gen_server:call(find_pool(Pool), stats). %% @doc start a pool start_pool(Name, Options) -> case find_pool(Name, Options) of Pid when is_pid(Pid) -> ok; Error -> Error end. %% @doc stop a pool stop_pool(Name) -> case find_pool(Name) of undefined -> ok; _Pid -> case supervisor:terminate_child(hackney_sup, Name) of ok -> _= supervisor:delete_child(hackney_sup, Name), ets:delete(hackney_pool, Name), ok; Error -> Error end end. notify(Pool, Msg) -> case find_pool(Pool) of undefined -> ok; Pid -> Pid ! Msg end. %% %% util functions for this pool %% %% @doc return a child spec suitable for embeding your pool in the %% supervisor child_spec(Name, Options0) -> Options = [{name, Name} | Options0], {Name, {hackney_pool, start_link, [Name, Options]}, permanent, 10000, worker, [hackney_pool]}. %% @doc get the number of connections in the pool count(Name) -> gen_server:call(find_pool(Name), count). %% @doc get the number of connections in the pool for `{Host0, Port, Transport}' count(Name, {Host0, Port, Transport}) -> Host = string:to_lower(Host0), gen_server:call(find_pool(Name), {count, {Host, Port, Transport}}). %% @doc get max pool size max_connections(Name) -> gen_server:call(find_pool(Name), max_connections). %% @doc change the pool size set_max_connections(Name, NewSize) -> gen_server:cast(find_pool(Name), {set_maxconn, NewSize}). %% @doc get timeout timeout(Name) -> gen_server:call(find_pool(Name), timeout). %% @doc change the connection timeout %% set_timeout(Name, NewTimeout) -> gen_server:cast(find_pool(Name), {set_timeout, NewTimeout}). %% @private %% %% do_start_pool(Name, Options) -> Spec = child_spec(Name, Options), case supervisor:start_child(hackney_sup, Spec) of {ok, Pid} -> Pid; {error, {already_started, _}} -> find_pool(Name, Options) end. find_pool(Name) -> case ets:lookup(?MODULE, Name) of [] -> undefined; [{_, Pid}] -> Pid end. find_pool(Name, Options) -> case ets:lookup(?MODULE, Name) of [] -> do_start_pool(Name, Options); [{_, Pid}] -> Pid end. start_link(Name, Options0) -> Options = hackney_util:maybe_apply_defaults([max_connections, timeout], Options0), gen_server:start_link(?MODULE, [Name, Options], []). init([Name, Options]) -> case lists:member({seed,1}, ssl:module_info(exports)) of true -> % Make sure that the ssl random number generator is seeded % This was new in R13 (ssl-3.10.1 in R13B vs. ssl-3.10.0 in R12B-5) apply(ssl, seed, [crypto:strong_rand_bytes(255)]); false -> ok end, MaxConn = case proplists:get_value(pool_size, Options) of undefined -> proplists:get_value(max_connections, Options); Size -> Size end, Timeout = proplists:get_value(timeout, Options), %% register the module ets:insert(?MODULE, {Name, self()}), %% initialize metrics Engine = init_metrics(Name), {ok, #state{name=Name, metrics=Engine, max_connections=MaxConn, timeout=Timeout}}. handle_call(stats, _From, State) -> {reply, handle_stats(State), State}; handle_call(count, _From, #state{sockets=Sockets}=State) -> {reply, dict:size(Sockets), State}; handle_call(timeout, _From, #state{timeout=Timeout}=State) -> {reply, Timeout, State}; handle_call(max_connections, _From, #state{max_connections=MaxConn}=State) -> {reply, MaxConn, State}; handle_call({checkout, Dest, Pid, RequestRef}, From, State) -> #state{name=PoolName, metrics = Engine, max_connections=MaxConn, clients=Clients, queues = Queues, pending = Pending} = State, {Reply, State2} = find_connection(Dest, Pid, State), case Reply of {ok, _Socket, _Owner} -> State3 = monitor_client(Dest, RequestRef, State2), ok = update_usage(State3), {reply, Reply, State3}; no_socket -> case dict:size(Clients) >= MaxConn of true -> Queues2 = add_to_queue(Dest, From, RequestRef, Queues), Pending2 =add_pending(RequestRef, From, Dest, Pending), _ = metrics:update_histogram( Engine, [hackney_pool, PoolName, queue_count], dict:size(Pending2) ), {noreply, State2#state{queues = Queues2, pending = Pending2}}; false -> State3 = monitor_client(Dest, RequestRef, State2), ok = update_usage(State3), {reply, {error, no_socket, self()}, State3} end end; handle_call({checkin, Ref, Dest, Socket, Transport}, From, State) -> gen_server:reply(From, ok), Clients2 = case dict:find(Ref, State#state.clients) of {ok, Dest} -> dict:erase(Ref, State#state.clients); error -> State#state.clients end, State2 = case Transport:peername(Socket) of {ok, {_Adress, _Port}} -> %% socket is not closed, try to deliver it or store it deliver_socket(Socket, Dest, State#state{clients=Clients2}); Error -> %% socket may be half-closed, close it and return catch Transport:close(Socket), ?report_trace("checkin: socket is not ok~n", [{socket, Socket}, {peername, Error}]), State#state{clients=Clients2} end, ok = update_usage(State2), {noreply, State2}; handle_call({count, Key}, _From, #state{connections=Conns}=State) -> Size = case dict:find(Key, Conns) of {ok, Sockets} -> length(Sockets); error -> 0 end, {reply, Size, State}. handle_cast({set_maxconn, MaxConn}, State) -> {noreply, State#state{max_connections=MaxConn}}; handle_cast({set_timeout, NewTimeout}, State) -> {noreply, State#state{timeout=NewTimeout}}; handle_cast({checkout_cancel, Dest, Ref}, State) -> #state{queues=Queues, pending=Pending} = State, {Queues2, Removed} = del_from_queue(Dest, Ref, Queues), case Removed of true -> Pending2 = del_pending(Ref, Pending), {noreply, State#state{queues=Queues2, nb_waiters=Pending2}}; false -> % we leak the socket here but 'DOWN' will mop up for us when it times out {noreply, dequeue(Dest, Ref, State)} end; handle_cast(_Msg, State) -> {noreply, State}. handle_info({timeout, Socket}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({tcp, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({tcp_closed, Socket}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({ssl, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({ssl_closed, Socket}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({tcp_error, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({ssl_error, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({'DOWN', Ref, request, _Pid, _Reason}, State) -> case dict:find(Ref, State#state.clients) of {ok, Dest} -> {noreply, dequeue(Dest, Ref, State)}; error -> NewState = remove_pending(Ref, State), {noreply, NewState} end; handle_info(_, State) -> {noreply, State}. code_change(_OldVsn, State, _Extra) -> {ok, State}. terminate(_Reason, #state{name=PoolName, metrics=Engine, sockets=Sockets}) -> %% close any sockets in the pool lists:foreach(fun({Socket, {{_, _, Transport}, Timer}}) -> cancel_timer(Socket, Timer), Transport:close(Socket) end, dict:to_list(Sockets)), %% delete pool metrics ok = delete_metrics(Engine, PoolName), ok. %% internals dequeue(Dest, Ref, State) -> #state{clients=Clients, queues=Queues, pending=Pending} = State, Clients2 = dict:erase(Ref, Clients), case queue_out(Dest, Queues) of empty -> State#state{clients = Clients2}; {ok, {From, Ref2}, Queues2} -> Pending2 = del_pending(Ref, Pending), _ = metrics:update_histogram( State#state.metrics, [hackney_pool, State#state.name, queue_count], dict:size(Pending2) ), gen_server:reply(From, {error, no_socket, self()}), State2 = State#state{queues = Queues2, clients = Clients2, pending=Pending2}, monitor_client(Dest, Ref2, State2) end. find_connection({_Host, _Port, Transport}=Dest, Pid, #state{connections=Conns, sockets=Sockets}=State) -> case dict:find(Dest, Conns) of {ok, [S | Rest]} -> Transport:setopts(S, [{active, false}]), case sync_socket(Transport, S) of true -> case Transport:controlling_process(S, Pid) of ok -> {_, Timer} = dict:fetch(S, Sockets), cancel_timer(S, Timer), NewConns = update_connections(Rest, Dest, Conns), NewSockets = dict:erase(S, Sockets), NewState = State#state{connections=NewConns, sockets=NewSockets}, {{ok, S, self()}, NewState}; {error, badarg} -> %% something happened here normally the PID died, %% but make sure we still have the control of the %% process catch Transport:controlling_process(S, self()), %% and then close it find_connection(Dest, Pid, remove_socket(S, State)); _Else -> find_connection(Dest, Pid, remove_socket(S, State)) end; false -> ?report_trace("checkout: socket unsynced~n", []), find_connection(Dest, Pid, remove_socket(S, State)) end; _Else -> {no_socket, State} end. remove_socket(Socket, #state{connections=Conns, sockets=Sockets}=State) -> _ = metrics:update_histogram(State#state.metrics, [hackney, State#state.name, free_count], dict:size(Sockets)), case dict:find(Socket, Sockets) of {ok, {{_Host, _Port, Transport}=Key, Timer}} -> cancel_timer(Socket, Timer), catch Transport:close(Socket), ConnSockets = lists:delete(Socket, dict:fetch(Key, Conns)), NewConns = update_connections(ConnSockets, Key, Conns), NewSockets = dict:erase(Socket, Sockets), State#state{connections=NewConns, sockets=NewSockets}; error -> State end. store_socket({_Host, _Port, Transport} = Dest, Socket, #state{timeout=Timeout, connections=Conns, sockets=Sockets}=State) -> Timer = erlang:send_after(Timeout, self(), {timeout, Socket}), %% make sure to close the socket if anything is received while we are in %% the pool. Transport:setopts(Socket, [{active, once}, {packet, 0}]), ConnSockets = case dict:find(Dest, Conns) of {ok, OldSockets} -> [Socket | OldSockets]; error -> [Socket] end, State#state{connections = dict:store(Dest, ConnSockets, Conns), sockets = dict:store(Socket, {Dest, Timer}, Sockets)}. update_connections([], Key, Connections) -> dict:erase(Key, Connections); update_connections(Sockets, Key, Connections) -> dict:store(Key, Sockets, Connections). cancel_timer(Socket, Timer) -> case erlang:cancel_timer(Timer) of false -> receive {timeout, Socket} -> ok after 0 -> ok end; _ -> ok end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ add_to_queue({_Host, _Port, _Transport} = Dest, From, Ref, Queues) -> case dict:find(Dest, Queues) of error -> dict:store(Dest, queue:in({From, Ref}, queue:new()), Queues); {ok, Q} -> dict:store(Dest, queue:in({From, Ref}, Q), Queues) end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ del_from_queue({_Host, _Port, _Transport} = Dest, Ref, Queues) -> case dict:find(Dest, Queues) of error -> {Queues, false}; {ok, Q} -> Q2 = queue:filter(fun({_, R}) -> R =/= Ref end, Q), Removed = queue:len(Q) =/= queue:len(Q2), Queues2 = case queue:is_empty(Q2) of true -> dict:erase(Dest, Queues); false -> dict:store(Dest, Q2, Queues) end, {Queues2, Removed} end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ queue_out({_Host, _Port, _Transport} = Dest, Queues) -> case dict:find(Dest, Queues) of error -> empty; {ok, Q} -> {{value, {From, Ref}}, Q2} = queue:out(Q), Queues2 = case queue:is_empty(Q2) of true -> dict:erase(Dest, Queues); false -> dict:store(Dest, Q2, Queues) end, {ok, {From, Ref}, Queues2} end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ deliver_socket(Socket, {_, _, Transport} = Dest, State) -> #state{queues = Queues, pending=Pending} = State, case queue_out(Dest, Queues) of empty -> store_socket(Dest, Socket, State); {ok, {{PidWaiter, _} = FromWaiter, Ref}, Queues2} -> Pending2 = del_pending(Ref, Pending), _ = metrics:update_histogram( State#state.metrics, [hackney_pool, State#state.name, queue_count], dict:size(Pending2) ), case Transport:controlling_process(Socket, PidWaiter) of ok -> gen_server:reply(FromWaiter, {ok, Socket, self()}), monitor_client(Dest, Ref, State#state{queues = Queues2, pending=Pending2}); _Error -> % Something wrong, close the socket _ = (catch Transport:close(Socket)), %% and let the waiter connect to a new one gen_server:reply(FromWaiter, {error, no_socket, self()}), State#state{queues = Queues2, pending = Pending2} end end. %% check that no events from the sockets is received after setting it to %% passive. sync_socket(Transport, Socket) -> {Msg, MsgClosed, MsgError} = Transport:messages(Socket), receive {Msg, Socket, _} -> false; {MsgClosed, Socket} -> false; {MsgError, Socket, _} -> false after 0 -> true end. add_pending(Ref, From, Dest, Pending) -> dict:store(Ref, {From, Dest}, Pending). del_pending(Ref, Pending) -> dict:erase(Ref, Pending). remove_pending(Ref, #state{queues=Queues0, pending=Pending0} = State) -> case dict:find(Ref, Pending0) of {ok, {From, Dest}} -> Pending1 = dict:erase(Ref, Pending0), Queues1 = case dict:find(Dest, Queues0) of {ok, Q0} -> Q1 = queue:filter( fun (PendingReq) when PendingReq =:= {From, Ref} -> false; (_) -> true end, Q0 ), dict:store(Dest, Q1, Queues0); error -> Queues0 end, State#state{queues=Queues1, pending=Pending1}; error -> State end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ monitor_client(Dest, Ref, State) -> Clients2 = dict:store(Ref, Dest, State#state.clients), State#state{clients = Clients2}. init_metrics(PoolName) -> %% get metrics module Engine = hackney_metrics:get_engine(), %% initialise metrics _ = metrics:new(Engine, histogram, [hackney_pool, PoolName, take_rate]), _ = metrics:new(Engine, counter, [hackney_pool, PoolName, no_socket]), _ = metrics:new(Engine, histogram, [hackney_pool, PoolName, in_use_count]), _ = metrics:new(Engine, histogram, [hackney_pool, PoolName, free_count]), _ = metrics:new(Engine, histogram, [hackney_pool, PoolName, queue_counter]), Engine. delete_metrics(Engine, PoolName) -> _ = metrics:delete(Engine, [hackney_pool, PoolName, take_rate]), _ = metrics:delete(Engine, [hackney_pool, PoolName, no_socket]), _ = metrics:delete(Engine, [hackney_pool, PoolName, in_use_count]), _ = metrics:delete(Engine, [hackney_pool, PoolName, free_count]), _ = metrics:delete(Engine, [hackney_pool, PoolName, queue_counter]), ok. update_usage( #state{name=PoolName, metrics=Engine, sockets=Sockets, clients=Clients} ) -> _ = metrics:update_histogram(Engine, [hackney_pool, PoolName,in_use_count], dict:size(Clients) - 1), _ = metrics:update_histogram(Engine, [hackney_pool, PoolName, free_count], dict:size(Sockets) - 1), ok. handle_stats(State) -> #state{name=PoolName, max_connections=Max, sockets=Sockets, clients=Clients, pending=Pending} = State, [{name, PoolName}, {max, Max}, {in_use_count, dict:size(Clients)}, {free_count, dict:size(Sockets)}, {queue_count, dict:size(Pending)}].