fix: pool http+unix connections over the unix socket

src/hackney.erl

@@ -400,10 +400,9 @@ connect_pool_new(Transport, Host, Port, Options, FinalSslOpts, PoolHandler) ->
           %% an exact match of the TLS options hash (tls_key in Options)
           connect_pool_ssl(Transport, Host, Port, Options, FinalSslOpts, PoolHandler);
         false ->
-          %% Always checkout as TCP - pool only stores TCP connections
-          case PoolHandler:checkout(Host, Port, hackney_tcp, Options) of
+          case PoolHandler:checkout(Host, Port, pool_transport(Transport), Options) of
             {ok, PoolRef, ConnPid} ->
-              %% Got TCP connection - upgrade to SSL if needed
+              %% Upgrade to SSL if this is an HTTPS request
               case maybe_upgrade_ssl(Transport, ConnPid, FinalSslOpts) of
                 ok ->
                   %% Check if HTTP/2 was negotiated, register for multiplexing
@@ -424,6 +423,12 @@ connect_pool_new(Transport, Host, Port, Options, FinalSslOpts, PoolHandler) ->
       {error, checkout_timeout}
   end.
 
+%% @private Transport the pool dials for a request. The pool does not store
+%% TLS sockets: HTTPS dials TCP and is upgraded in place. Every other
+%% transport, such as a unix domain socket, is dialled as itself.
+pool_transport(hackney_ssl) -> hackney_tcp;
+pool_transport(Transport) -> Transport.
+
 %% @private SSL-pooling checkout. A `ready' conn is an already-upgraded
 %% HTTPS/1.1 connection reused on an exact tls_key match; it was registered
 %% for h2 at creation if applicable, so it is not re-registered here. A

src/hackney_pool.erl

@@ -729,11 +729,11 @@ handle_cast({set_timeout, NewTimeout}, State) ->
     {noreply, State#state{keepalive_timeout=Capped}};
 
 handle_cast({prewarm, Host, Port}, #state{prewarm_count=Count}=State) ->
-    State2 = do_prewarm(Host, Port, Count, State),
+    State2 = do_prewarm(connection_key(Host, Port, hackney_tcp), Count, State),
     {noreply, State2};
 
 handle_cast({prewarm, Host, Port, Count}, State) ->
-    State2 = do_prewarm(Host, Port, Count, State),
+    State2 = do_prewarm(connection_key(Host, Port, hackney_tcp), Count, State),
     {noreply, State2};
 
 handle_cast({prewarm_checkin, Pid, Key}, State) ->
@@ -950,8 +950,12 @@ connection_key(Host, Port, Transport) ->
     connection_key(Host, Port, Transport, default).
 
 connection_key(Host0, Port, Transport, TlsKey) ->
-    Host = string:lowercase(Host0),
-    {Host, Port, Transport, TlsKey}.
+    {normalize_host(Host0, Transport), Port, Transport, TlsKey}.
+
+%% @private The pool dials the host stored in the key. Hostnames are
+%% case-insensitive; a unix domain socket path is not.
+normalize_host(Path, hackney_local_tcp) -> Path;
+normalize_host(Host, _Transport) -> string:lowercase(Host).
 
 %% @private Host and port of a pool connection key.
 key_host_port({Host, Port, _Transport, _TlsKey}) ->
@@ -1265,10 +1269,11 @@ keepalive_ready(Info) ->
         maps:get(ready, Info, false) =:= true.
 
 %% @private Close branch of a checkin: drop the monitor and keep the host's
-%% TCP prewarm warm (the replacement for a closed conn is a TCP conn that
-%% gets upgraded on its next SSL checkout). The conn itself is stopped by
-%% the caller: stop_conn for async checkin, a stop cast for sync checkin
-%% where the conn is blocked waiting on the pool's reply.
+%% prewarm warm on the transport the conn was dialled with (the replacement
+%% for a closed HTTPS conn is a TCP conn that gets upgraded on its next SSL
+%% checkout). The conn itself is stopped by the caller: stop_conn for async
+%% checkin, a stop cast for sync checkin where the conn is blocked waiting on
+%% the pool's reply.
 checkin_close(Pid, Key, InUse2, State) ->
     #state{pid_monitors=PidMonitors} = State,
     PidMonitors2 = case maps:take(Pid, PidMonitors) of
@@ -1278,9 +1283,15 @@ checkin_close(Pid, Key, InUse2, State) ->
         error -> PidMonitors
     end,
     State2 = State#state{in_use=InUse2, pid_monitors=PidMonitors2},
-    {Host, Port} = key_host_port(Key),
-    TcpKey = connection_key(Host, Port, hackney_tcp),
-    maybe_maintain_prewarm(TcpKey, activate_host(TcpKey, State2)).
+    DialKey = dial_key(Key),
+    maybe_maintain_prewarm(DialKey, activate_host(DialKey, State2)).
+
+%% @private Key of the plain bucket a conn under Key was dialled into: an
+%% HTTPS conn was dialled as TCP, any other conn as its own transport.
+dial_key({Host, Port, hackney_ssl, _TlsKey}) ->
+    connection_key(Host, Port, hackney_tcp);
+dial_key({Host, Port, Transport, _TlsKey}) ->
+    connection_key(Host, Port, Transport).
 
 %% @private Pool branch of a checkin: hand ownership to the pool and store
 %% the conn under its checkout key (TCP or SSL bucket).
@@ -1429,11 +1440,12 @@ do_unregister_h3(Pid, State) ->
     State#state{h3_connections = H3Conns2, pid_monitors = PidMonitors2}.
 
 %% @private Prewarm connections to a host
-%% Creates connections up to Count if there are fewer available.
-%% Only creates TCP connections (prewarm doesn't apply to SSL).
-do_prewarm(Host, Port, Count, State) ->
+%% Creates connections up to Count if there are fewer available, dialled
+%% with the transport of Key. Key is a plain bucket key: an HTTPS host is
+%% prewarmed through its TCP key (prewarm doesn't apply to SSL).
+do_prewarm(Key, Count, State) ->
     #state{name=PoolName, available=Available, keepalive_timeout=IdleTimeout} = State,
-    Key = connection_key(Host, Port, hackney_tcp),
+    {Host, Port} = key_host_port(Key),
     CurrentCount = case maps:find(Key, Available) of
         {ok, Pids} -> length(Pids);
         error -> 0
@@ -1448,7 +1460,7 @@ do_prewarm(Host, Port, Count, State) ->
             %% Create connections in background to not block the pool
             PoolPid = self(),
             spawn(fun() ->
-                prewarm_connections(PoolPid, Host, Port, Needed, IdleTimeout)
+                prewarm_connections(PoolPid, Key, Needed, IdleTimeout)
             end),
             State2;
         false ->
@@ -1474,8 +1486,7 @@ maybe_maintain_prewarm(Key, State) ->
             end,
             case CurrentCount < PrewarmCount of
                 true ->
-                    {Host, Port} = key_host_port(Key),
-                    do_prewarm(Host, Port, PrewarmCount, State);
+                    do_prewarm(Key, PrewarmCount, State);
                 false ->
                     State
             end;
@@ -1484,13 +1495,14 @@ maybe_maintain_prewarm(Key, State) ->
     end.
 
 %% @private Create prewarm connections (runs in separate process)
-prewarm_connections(_PoolPid, _Host, _Port, 0, _IdleTimeout) ->
+prewarm_connections(_PoolPid, _Key, 0, _IdleTimeout) ->
     ok;
-prewarm_connections(PoolPid, Host, Port, Count, IdleTimeout) ->
+prewarm_connections(PoolPid, Key, Count, IdleTimeout) ->
+    {Host, Port, Transport, _TlsKey} = Key,
     ConnOpts = #{
         host => Host,
         port => Port,
-        transport => hackney_tcp,
+        transport => Transport,
         connect_timeout => 5000,
         recv_timeout => 5000,
         idle_timeout => IdleTimeout,
@@ -1504,15 +1516,14 @@ prewarm_connections(PoolPid, Host, Port, Count, IdleTimeout) ->
             case connect_connection(Pid, ?PREWARM_CONNECT_TIMEOUT) of
                 ok ->
                     %% Checkin the new connection to the pool
-                    gen_server:cast(PoolPid, {prewarm_checkin, Pid,
-                                              connection_key(Host, Port, hackney_tcp)});
+                    gen_server:cast(PoolPid, {prewarm_checkin, Pid, Key});
                 {error, _Reason} ->
                     stop_conn(Pid)
             end;
         {error, _Reason} ->
             ok
     end,
-    prewarm_connections(PoolPid, Host, Port, Count - 1, IdleTimeout).
+    prewarm_connections(PoolPid, Key, Count - 1, IdleTimeout).
 
 handle_stats(State) ->
     #state{name=PoolName, max_connections=Max, available=Available,

test/hackney_pool_unix_socket_tests.erl

@@ -0,0 +1,94 @@
+%% @doc Pooled requests over unix domain sockets (http+unix URLs).
+-module(hackney_pool_unix_socket_tests).
+-include_lib("eunit/include/eunit.hrl").
+
+-define(LISTENER, test_pool_unix_socket_http).
+-define(POOL, test_pool_unix_socket).
+-define(PREWARM_POOL, test_pool_unix_socket_prewarm).
+
+setup() ->
+    {ok, _} = application:ensure_all_started(hackney),
+    {ok, _} = application:ensure_all_started(cowboy),
+    %% sun_path is limited to about 100 bytes, so stay under /tmp rather than
+    %% a possibly long per-user temp dir. The mixed-case file name checks the
+    %% pool does not case-fold the path like a hostname.
+    Dir = filename:join("/tmp", "hackney-unix-" ++
+                                integer_to_list(erlang:unique_integer([positive]))),
+    ok = file:make_dir(Dir),
+    Path = filename:join(Dir, "Hackney.sock"),
+    Dispatch = cowboy_router:compile([{'_', [{"/delay/:seconds", delay_handler, []},
+                                             {"/[...]", test_http_resource, []}]}]),
+    {ok, _} = cowboy:start_clear(?LISTENER, [{ip, {local, Path}}, {port, 0}],
+                                 #{env => #{dispatch => Dispatch}}),
+    ok = hackney_pool:start_pool(?POOL, [{pool_size, 4}, {prewarm_count, 0}]),
+    ok = hackney_pool:start_pool(?PREWARM_POOL, [{pool_size, 1}, {prewarm_count, 2}]),
+    {Dir, Path}.
+
+cleanup({Dir, Path}) ->
+    hackney_pool:stop_pool(?PREWARM_POOL),
+    hackney_pool:stop_pool(?POOL),
+    cowboy:stop_listener(?LISTENER),
+    _ = file:delete(Path),
+    _ = file:del_dir(Dir),
+    ok.
+
+url(SocketPath, ReqPath) ->
+    Netloc = hackney_url:urlencode(list_to_binary(SocketPath)),
+    <<"http+unix://", Netloc/binary, ReqPath/binary>>.
+
+%% The default Host for an http+unix URL is the percent-encoded socket path,
+%% which cowboy rejects (a reg-name without pct-encoding), so name a host.
+unix_get(Path, ReqPath, Pool) ->
+    hackney:request(get, url(Path, ReqPath), [{<<"Host">>, <<"localhost">>}], <<>>,
+                    [{pool, Pool}, with_body]).
+
+%% The tests run in order: the first one relies on being the only client
+%% of the listener so far.
+pool_unix_socket_test_() ->
+    {setup,
+     fun setup/0,
+     fun cleanup/1,
+     fun({_Dir, Path}) ->
+         [
+          {"pooled requests over a unix socket reuse one connection",
+           fun() -> test_pooled_reuse(Path) end},
+          {"closing an overflow unix socket connection prewarms over the socket",
+           fun() -> test_prewarm_after_close(Path) end}
+         ]
+     end}.
+
+test_pooled_reuse(Path) ->
+    ?assertMatch({ok, 200, _, _}, unix_get(Path, <<"/get">>, ?POOL)),
+    ok = wait_until(fun() -> free_count(?POOL) =:= 1 end),
+    ?assertMatch({ok, 200, _, _}, unix_get(Path, <<"/get">>, ?POOL)),
+    ok = wait_until(fun() -> free_count(?POOL) =:= 1 end),
+    ?assertEqual(1, length(ranch:procs(?LISTENER, connections))).
+
+%% With pool_size 1, the second of two overlapping requests gets an overflow
+%% connection that is closed at checkin, which tops the host up to
+%% prewarm_count. The replacement must be dialled over the socket too.
+test_prewarm_after_close(Path) ->
+    Self = self(),
+    Callers = [spawn_link(fun() ->
+                              Self ! {self(), unix_get(Path, <<"/delay/1">>, ?PREWARM_POOL)}
+                          end)
+               || _ <- [1, 2]],
+    [?assertMatch({ok, 200, _, _}, receive {Caller, Res} -> Res after 5000 -> timeout end)
+     || Caller <- Callers],
+    ?assertEqual(ok, wait_until(fun() -> free_count(?PREWARM_POOL) =:= 2 end)).
+
+free_count(Pool) ->
+    proplists:get_value(free_count, hackney_pool:get_stats(Pool)).
+
+wait_until(Fun) ->
+    wait_until(Fun, 200).
+
+wait_until(_Fun, 0) ->
+    timeout;
+wait_until(Fun, Retries) ->
+    case Fun() of
+        true -> ok;
+        false ->
+            timer:sleep(10),
+            wait_until(Fun, Retries - 1)
+    end.