hackney

%%% -*- 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 <benoitc@e-engura.org>

%% @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,
         find_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, [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]) ->
    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}.