From cf15810d709028421cd836272f14aa50f9058091 Mon Sep 17 00:00:00 2001
From: leonardb <leonardb@users.noreply.github.com>
Date: Mon, 10 Sep 2018 09:45:25 -0400
Subject: [PATCH] (feat) Support 'checkout_timeout' option (#520)

* (feat) Support 'checkout_timeout' option

Allow caller to handle overload when pool is maxed out rather than returning
incorrect 'connection_timeout' response

* (fix) fix timeout on pool test causing CI failure

* (fix) Fall back to connect_timeout when checkout_timeout is not set
---
 src/hackney.erl             |  2 ++
 src/hackney_pool.erl        |  6 ++++--
 test/hackney_pool_tests.erl | 18 ++++++++++++++++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/hackney.erl b/src/hackney.erl
index a8482ba5..111cab5b 100644
--- a/src/hackney.erl
+++ b/src/hackney.erl
@@ -238,6 +238,8 @@ request(Method, URL, Headers, Body) ->
 %%          <li>`{proxy, proxy_options()}': to connect via a proxy.</li>
 %%          <li>`insecure': to perform "insecure" SSL connections and
 %%          transfers without checking the certificate</li>
+%%          <li>`{checkout_timeout, infinity | integer()}': timeout used when
+%%          checking out a socket from the pool, in milliseconds. Default is 8000</li>
 %%          <li>`{connect_timeout, infinity | integer()}': timeout used when
 %%          establishing a connection, in milliseconds. Default is 8000</li>
 %%          <li>`{recv_timeout, infinity | integer()}': timeout used when
diff --git a/src/hackney_pool.erl b/src/hackney_pool.erl
index fdabbfb0..599fcf5c 100644
--- a/src/hackney_pool.erl
+++ b/src/hackney_pool.erl
@@ -68,8 +68,10 @@ checkout(Host0, Port, Transport, #client{options=Opts}=Client) ->
   Name = proplists:get_value(pool, Opts, default),
   Pool = find_pool(Name, Opts),
   ConnectTimeout = proplists:get_value(connect_timeout, Opts, 8000),
+  %% Fall back to using connect_timeout if checkout_timeout is not set
+  CheckoutTimeout = proplists:get_value(checkout_timeout, Opts, ConnectTimeout),
   case catch gen_server:call(Pool, {checkout, {Host, Port, Transport},
-    Pid, RequestRef}, ConnectTimeout) of
+    Pid, RequestRef}, CheckoutTimeout) of
     {ok, Socket, Owner} ->
       CheckinReference = {Host, Port, Transport},
       {ok, {Name, RequestRef, CheckinReference, Owner, Transport}, Socket};
@@ -82,7 +84,7 @@ checkout(Host0, Port, Transport, #client{options=Opts}=Client) ->
     {'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}
+      {error, checkout_timeout}
   end.
 
 %% @doc release a socket in the pool
diff --git a/test/hackney_pool_tests.erl b/test/hackney_pool_tests.erl
index 5896835f..4df20a3c 100644
--- a/test/hackney_pool_tests.erl
+++ b/test/hackney_pool_tests.erl
@@ -8,7 +8,8 @@ dummy_test() ->
 
 multipart_test_() ->
     {setup, fun start/0, fun stop/1,
-      [queue_timeout()]}.
+      [{timeout, 120, queue_timeout()},
+       {timeout, 120, checkout_timeout()}]}.
 
 start() ->
     error_logger:tty(false),
@@ -39,7 +40,20 @@ queue_timeout() ->
                 ok = hackney:finish_send_body(Ref),
                 {ok, _Status, _Headers, Ref} = hackney:start_response(Ref),
                 ok = hackney:skip_body(Ref),
-                {ok, _} = hackney:request(post, URL, Headers, stream, Opts)
+                {ok, Ref2} = hackney:request(post, URL, Headers, stream, Opts),
+                hackney:close(Ref2)
         end
     end.
 
+checkout_timeout() ->
+    fun() ->
+        URL = <<"http://localhost:8123/pool">>,
+        Headers = [],
+        Opts = [{max_body, 2048}, {pool, pool_test}, {connect_timeout, 1000}, {checkout_timeout, 100}],
+        case hackney:request(post, URL, Headers, stream, Opts) of
+            {ok, Ref} ->
+                {error, Error} = hackney:request(post, URL, Headers, stream, Opts),
+                hackney:close(Ref),
+                ?assertEqual(Error, checkout_timeout)
+        end
+    end.