%%% -*- 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-2014, Benoît Chesneau %% @doc pool of sockets connections %% -module(hackney_pool). -behaviour(gen_server). -behaviour(hackney_pool_handler). %% PUBLIC API -export([start/0, checkout/4, checkin/2]). -export([start_pool/2, stop_pool/1]). -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"). -record(state, { max_connections, timeout, clients = dict:new(), queues = dict:new(), % Dest => queue of Froms connections = dict:new(), sockets = dict:new()}). 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}) -> Host = string:to_lower(Host0), Pid = self(), Name = proplists:get_value(pool, Opts, default), Pool = find_pool(Name, Opts), case gen_server:call(Pool, {checkout, {Host, Port, Transport}, Pid}, infinity) of {ok, Socket, Owner} -> CheckinReference = {Host, Port, Transport}, {ok, {Name, CheckinReference, Owner, Transport}, Socket}; {error, no_socket, Owner} -> CheckinReference = {Host, Port, Transport}, {error, no_socket, {Name, CheckinReference, Owner, Transport}}; {error, Reason} -> {error, Reason} end. %% @doc release a socket in the pool checkin({_Name, Dest, Owner, Transport}, Socket) -> case Transport:controlling_process(Socket, Owner) of ok -> gen_server:call(Owner, {checkin, Dest, Socket, Transport}, infinity); _Error -> catch Transport:close(Socket), ok end. %% @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. %% %% 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, [Name]}. %% @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]) -> process_flag(priority, high), 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: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()}), {ok, #state{max_connections=MaxConn, timeout=Timeout}}. 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}, From, State) -> #state{max_connections=MaxConn, clients=Clients, queues = Queues} = State, {Reply, State2} = find_connection(Dest, Pid, State), case Reply of {ok, _Socket, _Owner} -> State3 = monitor_client(Dest, From, State2), {reply, Reply, State3}; no_socket -> case dict:size(Clients) >= MaxConn of true -> Queues2 = add_to_queue(Dest, From, Queues), {noreply, State2#state{queues = Queues2}}; false -> State3 = monitor_client(Dest, From, State2), {reply, {error, no_socket, self()}, State3} end end; handle_call({checkin, Dest, Socket, _Transport}, {Pid, _} = From, State) -> gen_server:reply(From, ok), Clients2 = case dict:find(Pid, State#state.clients) of {ok, {Dest, MonRef}} -> true = erlang:demonitor(MonRef, [flush]), dict:erase(Pid, State#state.clients); error -> State#state.clients end, State2 = deliver_socket(Socket, Dest, State#state{clients=Clients2}), {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(_Msg, State) -> {noreply, State}. handle_info({timeout, Socket}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({tcp_closed, 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({tcp, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({ssl, Socket, _}, State) -> {noreply, remove_socket(Socket, State)}; handle_info({'DOWN', MonRef, process, Pid, _Reason}, State) -> {Dest, MonRef} = dict:fetch(Pid, State#state.clients), Clients2 = dict:erase(Pid, State#state.clients), case queue_out(Dest, State#state.queues) of empty -> {noreply, State#state{clients = Clients2}}; {ok, From, Queues2} -> gen_server:reply(From, {error, no_socket, self()}), State2 = State#state{queues = Queues2, clients = Clients2}, {noreply, monitor_client(Dest, From, State2)} end; handle_info(_, State) -> {noreply, State}. code_change(_OldVsn, State, _Extra) -> {ok, State}. terminate(_Reason, #state{sockets=Sockets}) -> lists:foreach(fun({Socket, {{_, _, Transport}, Timer}}) -> cancel_timer(Socket, Timer), Transport:close(Socket) end, dict:to_list(Sockets)), ok. %% internals 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 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} -> % Pid has timed out, reuse for someone else Transport:setopts(S, [{active, once}]), {no_socket, State}; _Else -> find_connection(Dest, Pid, remove_socket(S, State)) end; _Else -> {no_socket, State} end. remove_socket(Socket, #state{connections=Conns, sockets=Sockets}=State) -> 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}), ok = Transport:setopts(Socket, [{active, once}]), 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, Queues) -> case dict:find(Dest, Queues) of error -> dict:store(Dest, queue:in(From, queue:new()), Queues); {ok, Q} -> dict:store(Dest, queue:in(From, Q), Queues) end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ queue_out({_Host, _Port, _Transport} = Dest, Queues) -> case dict:find(Dest, Queues) of error -> empty; {ok, Q} -> {{value, From}, 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, Queues2} end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ deliver_socket(Socket, {_, _, Transport} = Dest, State) -> case queue_out(Dest, State#state.queues) of empty -> store_socket(Dest, Socket, State); {ok, {PidWaiter, _} = FromWaiter, Queues2} -> Transport:setopts(Socket, [{active, false}]), case Transport:controlling_process(Socket, PidWaiter) of ok -> gen_server:reply(FromWaiter, {ok, Socket, self()}), monitor_client(Dest, FromWaiter, State#state{queues = Queues2}); {error, badarg} -> % Pid died, reuse for someone else Transport:setopts(Socket, [{active, once}]), deliver_socket(Socket, Dest, State#state{queues = Queues2}); _Error -> % Something wrong with the socket; just remove it catch Transport:close(Socket), State end end. %------------------------------------------------------------------------------ %% @private %%------------------------------------------------------------------------------ monitor_client(Dest, {Pid, _} = _From, State) -> MonRef = erlang:monitor(process, Pid), Clients2 = dict:store(Pid, {Dest, MonRef}, State#state.clients), State#state{clients = Clients2}.