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-2024, Benoît Chesneau <benoitc@e-engura.org>

%% @doc Pool of connection processes.
%%
%% This module manages hackney_conn processes in a pool. Instead of storing
%% raw sockets, it tracks connection process pids. Idle timeout is handled
%% by the connection processes themselves (gen_statem state_timeout).
%%
-module(hackney_pool).
-behaviour(gen_server).

%% PUBLIC API
-export([start/0,
         checkout/4,
         checkin/2]).

%% HTTP/2 connection pooling
-export([checkout_h2/4,
         register_h2/5,
         unregister_h2/2,
         unregister_h2_all/0]).

%% HTTP/3 connection pooling
-export([checkout_h3/4,
         register_h3/5,
         unregister_h3/2,
         get_h3_session/4,
         store_h3_session/5,
         delete_h3_session/4]).

-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,
         prewarm/3,
         prewarm/4,
         host_stats/3,
         child_spec/2]).

-export([start_link/2]).

%% For internal use
-export([to_pool_name/1]).

%% 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,
    max_connections,
    keepalive_timeout,
    prewarm_count,
    %% Available connection processes: #{Key => [Pid]}
    available = #{},
    %% In-use connections: #{Pid => Key}
    in_use = #{},
    %% Pid to monitor ref mapping: #{Pid => MonitorRef}
    pid_monitors = #{},
    %% Hosts that have been activated (prewarm triggered): sets:set(Key)
    activated_hosts = sets:new(),
    %% HTTP/2 connections: #{Key => Pid} - one multiplexed connection per host
    %% These connections are shared across callers (not checked out exclusively)
    h2_connections = #{},
    %% HTTP/3 connections: #{Key => Pid} - one multiplexed QUIC connection per host
    %% These connections are shared across callers (not checked out exclusively)
    h3_connections = #{},
    %% HTTP/3 0-RTT/resumption session tickets: #{Key => Ticket} keyed by
    %% {Host, Port, Transport}. Replayed on the next connect to resume.
    h3_sessions = #{}
}).

-define(DEFAULT_MAX_CONNECTIONS, 50).
-define(DEFAULT_KEEPALIVE_TIMEOUT, 2000).  % 2 seconds max idle
-define(DEFAULT_PREWARM_COUNT, 4).         % Connections to maintain per host

start() ->
    %% Create ETS table to store pool pid by name
    _ = ets:new(?MODULE, [set, public, named_table, {read_concurrency, true}]),
    ok.

%% @doc Checkout a connection process from the pool.
%% Returns {ok, PoolInfo, Pid} where Pid is a hackney_conn process.
-spec checkout(Host :: string(), Port :: non_neg_integer(),
               Transport :: module(), Options :: list()) ->
    {ok, term(), pid()} | {error, term()}.
checkout(Host, Port, Transport, Options) ->
    Requester = self(),
    try
        do_checkout(Requester, Host, Port, Transport, Options)
    catch
        exit:{timeout, _} ->
            ?report_trace("pool: checkout timeout", []),
            {error, checkout_timeout};
        _:Error ->
            ?report_trace("pool: checkout failure", [{error, Error}]),
            {error, checkout_failure}
    end.

do_checkout(Requester, Host, Port, Transport, Opts) ->
    ConnectTimeout = proplists:get_value(connect_timeout, Opts, 8000),
    CheckoutTimeout = proplists:get_value(checkout_timeout, Opts, ConnectTimeout),
    PoolName = proplists:get_value(pool, Opts, default),
    Pool = find_pool(PoolName, Opts),
    Key = connection_key(Host, Port, Transport),

    case gen_server:call(Pool, {checkout, Key, Requester, Opts}, CheckoutTimeout) of
        {ok, Pid} ->
            %% Return pool info for later checkin
            PoolInfo = {PoolName, Key, Pool, Transport},
            {ok, PoolInfo, Pid};
        {error, _} = Error ->
            Error;
        {'EXIT', {timeout, _}} ->
            {error, checkout_timeout}
    end.

%% @doc Return a connection process to the pool.
-spec checkin(PoolInfo :: term(), Pid :: pid()) -> ok.
checkin({_PoolName, _Key, Pool, _Transport}, Pid) ->
    gen_server:cast(Pool, {checkin, nil, Pid}),
    ok.

%%====================================================================
%% HTTP/2 Connection Pooling
%%====================================================================

%% @doc Get an existing HTTP/2 connection for a host/port, or 'none' if not available.
%% HTTP/2 connections are shared (multiplexed) across callers.
-spec checkout_h2(Host :: string(), Port :: non_neg_integer(),
                  Transport :: module(), Options :: list()) ->
    {ok, pid()} | none.
checkout_h2(Host, Port, Transport, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    ConnectTimeout = proplists:get_value(connect_timeout, Options, 8000),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    try
        gen_server:call(Pool, {checkout_h2, Key}, ConnectTimeout)
    catch
        exit:{timeout, _} -> none;
        _:_ -> none
    end.

%% @doc Register an HTTP/2 connection in the pool for sharing.
%% Called after ALPN negotiation confirms HTTP/2.
-spec register_h2(Host :: string(), Port :: non_neg_integer(),
                  Transport :: module(), Pid :: pid(), Options :: list()) -> ok.
register_h2(Host, Port, Transport, Pid, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    gen_server:cast(Pool, {register_h2, Key, Pid}),
    ok.

%% @doc Remove an HTTP/2 connection from the pool (e.g., on GOAWAY).
-spec unregister_h2(Pid :: pid(), Options :: list()) -> ok.
unregister_h2(Pid, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    gen_server:cast(Pool, {unregister_h2, Pid}),
    ok.

%% @doc Remove all HTTP/2 connections from the default pool.
%% Used for testing to ensure clean state between tests.
-spec unregister_h2_all() -> ok.
unregister_h2_all() ->
    Pool = find_pool(default, []),
    gen_server:call(Pool, unregister_h2_all).

%%====================================================================
%% HTTP/3 Connection Pooling
%%====================================================================

%% @doc Get an existing HTTP/3 connection for a host/port, or 'none' if not available.
%% HTTP/3 connections are shared (multiplexed) across callers via QUIC streams.
-spec checkout_h3(Host :: string(), Port :: non_neg_integer(),
                  Transport :: module(), Options :: list()) ->
    {ok, pid()} | none.
checkout_h3(Host, Port, Transport, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    ConnectTimeout = proplists:get_value(connect_timeout, Options, 8000),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    try
        gen_server:call(Pool, {checkout_h3, Key}, ConnectTimeout)
    catch
        exit:{timeout, _} -> none;
        _:_ -> none
    end.

%% @doc Register an HTTP/3 connection in the pool for sharing.
%% Called after QUIC connection is established with HTTP/3.
-spec register_h3(Host :: string(), Port :: non_neg_integer(),
                  Transport :: module(), Pid :: pid(), Options :: list()) -> ok.
register_h3(Host, Port, Transport, Pid, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    gen_server:cast(Pool, {register_h3, Key, Pid}),
    ok.

%% @doc Remove an HTTP/3 connection from the pool (e.g., on connection close).
-spec unregister_h3(Pid :: pid(), Options :: list()) -> ok.
unregister_h3(Pid, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    gen_server:cast(Pool, {unregister_h3, Pid}),
    ok.

%% @doc Look up a cached HTTP/3 0-RTT/resumption session ticket for a host/port.
-spec get_h3_session(Host :: string(), Port :: non_neg_integer(),
                     Transport :: module(), Options :: list()) ->
    {ok, term()} | none.
get_h3_session(Host, Port, Transport, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    try
        gen_server:call(Pool, {get_h3_session, Key})
    catch
        _:_ -> none
    end.

%% @doc Cache an HTTP/3 session ticket for a host/port for later resumption.
-spec store_h3_session(Host :: string(), Port :: non_neg_integer(),
                       Transport :: module(), Ticket :: term(),
                       Options :: list()) -> ok.
store_h3_session(Host, Port, Transport, Ticket, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    gen_server:cast(Pool, {store_h3_session, Key, Ticket}),
    ok.

%% @doc Invalidate a cached HTTP/3 session ticket (e.g. after 0-RTT rejection).
-spec delete_h3_session(Host :: string(), Port :: non_neg_integer(),
                        Transport :: module(), Options :: list()) -> ok.
delete_h3_session(Host, Port, Transport, Options) ->
    PoolName = proplists:get_value(pool, Options, default),
    Pool = find_pool(PoolName, Options),
    Key = {Host, Port, Transport},
    gen_server:cast(Pool, {delete_h3_session, Key}),
    ok.

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.

%% @doc return a child spec suitable for embedding 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) ->
    case find_pool(Name) of
        undefined -> 0;
        Pid -> gen_server:call(Pid, count)
    end.

%% @doc get the number of connections in the pool for `{Host, Port, Transport}'
count(Name, Key) ->
    case find_pool(Name) of
        undefined -> 0;
        Pid -> gen_server:call(Pid, {count, Key})
    end.

%% @doc get max pool size
max_connections(Name) ->
    case find_pool(Name) of
        undefined -> 0;
        Pid -> gen_server:call(Pid, max_connections)
    end.

%% @doc change the pool size
set_max_connections(Name, NewSize) ->
    gen_server:cast(find_pool(Name), {set_maxconn, NewSize}).

%% @doc get timeout
timeout(Name) ->
    case find_pool(Name) of
        undefined -> 0;
        Pid -> gen_server:call(Pid, timeout)
    end.

%% @doc change the connection timeout
set_timeout(Name, NewTimeout) ->
    gen_server:cast(find_pool(Name), {set_timeout, NewTimeout}).

%% @doc Prewarm connections to a host (default count from pool settings)
%% Starts the pool if it doesn't exist.
-spec prewarm(atom(), string() | binary(), inet:port_number()) -> ok.
prewarm(PoolName, Host, Port) ->
    Pool = find_pool(PoolName, []),
    gen_server:cast(Pool, {prewarm, Host, Port}).

%% @doc Prewarm a specific number of connections to a host
%% Starts the pool if it doesn't exist.
-spec prewarm(atom(), string() | binary(), inet:port_number(), non_neg_integer()) -> ok.
prewarm(PoolName, Host, Port, Count) ->
    Pool = find_pool(PoolName, []),
    gen_server:cast(Pool, {prewarm, Host, Port, Count}).

%% @doc Get per-host connection statistics.
%% Returns a proplist with:
%% - active: number of active requests (from load_regulation)
%% - in_use: connections checked out from pool
%% - free: connections available in pool
-spec host_stats(atom(), string() | binary(), inet:port_number()) ->
    [{atom(), non_neg_integer()}].
host_stats(PoolName, Host, Port) ->
    Active = hackney_load_regulation:current(Host, Port),
    case find_pool(PoolName) of
        undefined ->
            [{active, Active}, {in_use, 0}, {free, 0}];
        Pool ->
            {InUse, Free} = gen_server:call(Pool, {host_stats, Host, Port}),
            [{active, Active}, {in_use, InUse}, {free, Free}]
    end.

to_pool_name(Name) when is_atom(Name) ->
    list_to_atom("hackney_pool_" ++ atom_to_list(Name));
to_pool_name(Name) when is_list(Name) ->
    list_to_atom("hackney_pool_" ++ Name);
to_pool_name(Name) when is_binary(Name) ->
    to_pool_name(binary_to_list(Name)).

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

%%====================================================================
%% gen_server callbacks
%%====================================================================

init([Name, Options]) ->
    MaxConn = case proplists:get_value(pool_size, Options) of
                  undefined ->
                      proplists:get_value(max_connections, Options, ?DEFAULT_MAX_CONNECTIONS);
                  Size ->
                      Size
              end,
    %% keepalive_timeout: max idle time for pooled connections (capped at 2s)
    %% Also accept 'timeout' for backward compatibility
    RawTimeout = case proplists:get_value(keepalive_timeout, Options) of
                     undefined ->
                         proplists:get_value(timeout, Options, ?DEFAULT_KEEPALIVE_TIMEOUT);
                     KT ->
                         KT
                 end,
    KeepaliveTimeout = min(RawTimeout, ?DEFAULT_KEEPALIVE_TIMEOUT),

    %% prewarm_count: number of TCP connections to maintain per host
    %% Check pool options first, then app env, then default
    PrewarmCount = case proplists:get_value(prewarm_count, Options) of
                       undefined ->
                           hackney_app:get_app_env(prewarm_count, ?DEFAULT_PREWARM_COUNT);
                       PC ->
                           PC
                   end,

    %% register the module
    ets:insert(?MODULE, {Name, self()}),

    {ok, #state{name=Name, max_connections=MaxConn,
                keepalive_timeout=KeepaliveTimeout, prewarm_count=PrewarmCount}}.

handle_call(stats, _From, State) ->
    {reply, handle_stats(State), State};

handle_call(count, _From, #state{available=Available, in_use=InUse}=State) ->
    AvailCount = maps:fold(fun(_, Pids, Acc) -> Acc + length(Pids) end, 0, Available),
    {reply, AvailCount + maps:size(InUse), State};

handle_call(timeout, _From, #state{keepalive_timeout=Timeout}=State) ->
    {reply, Timeout, State};

handle_call(max_connections, _From, #state{max_connections=MaxConn}=State) ->
    {reply, MaxConn, State};

handle_call({count, Key}, _From, #state{available=Available}=State) ->
    Count = case maps:find(Key, Available) of
                {ok, Pids} -> length(Pids);
                error -> 0
            end,
    {reply, Count, State};

handle_call({host_stats, Host, Port}, _From, #state{available=Available, in_use=InUse}=State) ->
    %% Count in_use and free for this host (any transport)
    HostLower = string:lowercase(Host),
    InUseCount = maps:fold(
        fun(_Pid, {H, P, _T}, Acc) when H =:= HostLower, P =:= Port -> Acc + 1;
           (_, _, Acc) -> Acc
        end, 0, InUse),
    FreeCount = maps:fold(
        fun({H, P, _T}, Pids, Acc) when H =:= HostLower, P =:= Port -> Acc + length(Pids);
           (_, _, Acc) -> Acc
        end, 0, Available),
    {reply, {InUseCount, FreeCount}, State};

handle_call({checkout, Key, Requester, Opts}, _From, State) ->
    #state{name=PoolName, max_connections=MaxConn,
           available=Available, in_use=InUse} = State,

    TotalInUse = maps:size(InUse),

    ?report_trace("pool: checkout request", [{pool, PoolName}, {key, Key},
        {total_in_use, TotalInUse}, {max_conn, MaxConn}]),

    case find_available(Key, Available) of
        {ok, Pid, Available2} ->
            %% Found an available connection - update owner to new requester
            ?report_debug("pool: reusing connection", [{pool, PoolName}, {pid, Pid}]),
            case hackney_conn:set_owner(Pid, Requester) of
                ok ->
                    InUse2 = maps:put(Pid, Key, InUse),
                    {reply, {ok, Pid}, State#state{available=Available2, in_use=InUse2}};
                {error, _} ->
                    %% #850: the connection closed between is_ready and
                    %% set_owner (server-side close race). It is already out of
                    %% Available2; drop it and start a fresh connection rather
                    %% than crashing the pool on a bad match.
                    case start_connection(Key, Requester, Opts, State#state{available=Available2}) of
                        {ok, Pid2, State2} ->
                            InUse2 = maps:put(Pid2, Key, State2#state.in_use),
                            {reply, {ok, Pid2}, State2#state{in_use=InUse2}};
                        {error, Reason} ->
                            {reply, {error, Reason}, State#state{available=Available2}}
                    end
            end;
        none when TotalInUse >= MaxConn ->
            %% At max connections - return error immediately (no queue)
            %% Note: In the new architecture, load_regulation handles waiting
            ?report_trace("pool: at max connections", [{pool, PoolName}, {in_use, TotalInUse}]),
            {reply, {error, checkout_timeout}, State};
        none ->
            %% No available connection, start a new one
            ?report_trace("pool: starting new connection", [{pool, PoolName}]),
            case start_connection(Key, Requester, Opts, State) of
                {ok, Pid, State2} ->
                    InUse2 = maps:put(Pid, Key, State2#state.in_use),
                    {reply, {ok, Pid}, State2#state{in_use=InUse2}};
                {error, Reason} ->
                    {reply, {error, Reason}, State}
            end
    end;

handle_call({checkin_sync, Pid}, _From, State) ->
    %% Synchronous checkin - caller waits for acknowledgement
    %% Legacy format without SSL flag - need to query connection (may deadlock if called from connection)
    State2 = do_checkin(Pid, State),
    {reply, ok, State2};

handle_call({checkin_sync, Pid, ShouldClose}, _From, State) ->
    %% Synchronous checkin with "should close" flag - avoids deadlock
    %% ShouldClose is true if connection was SSL upgraded or is a proxy tunnel
    State2 = do_checkin_with_close_flag(Pid, ShouldClose, State),
    {reply, ok, State2};

handle_call({checkout_h2, Key}, _From, #state{h2_connections = H2Conns} = State) ->
    %% HTTP/2 checkout - return existing connection if available.
    %% Liveness check includes both process_alive and gen_statem state, so a
    %% hackney_conn that already transitioned to `closed` (e.g. after an h2
    %% GOAWAY) but has not yet been removed via 'DOWN' is not handed out.
    case maps:get(Key, H2Conns, undefined) of
        undefined ->
            {reply, none, State};
        Pid ->
            case h2_conn_usable(Pid) of
                true ->
                    {reply, {ok, Pid}, State};
                false ->
                    H2Conns2 = maps:remove(Key, H2Conns),
                    {reply, none, State#state{h2_connections = H2Conns2}}
            end
    end;

handle_call({checkout_h3, Key}, _From, #state{h3_connections = H3Conns} = State) ->
    %% HTTP/3 checkout - return existing connection if available
    case maps:get(Key, H3Conns, undefined) of
        undefined ->
            {reply, none, State};
        Pid ->
            %% Verify connection is still alive
            case erlang:is_process_alive(Pid) of
                true ->
                    {reply, {ok, Pid}, State};
                false ->
                    %% Connection died, remove from pool
                    H3Conns2 = maps:remove(Key, H3Conns),
                    {reply, none, State#state{h3_connections = H3Conns2}}
            end
    end;

handle_call({get_h3_session, Key}, _From, #state{h3_sessions = Sessions} = State) ->
    case maps:get(Key, Sessions, undefined) of
        undefined -> {reply, none, State};
        Ticket -> {reply, {ok, Ticket}, State}
    end;

handle_call(unregister_h2_all, _From, State) ->
    %% Clear all HTTP/2 connections (for testing)
    {reply, ok, State#state{h2_connections = #{}}}.

handle_cast({checkin, _PoolInfo, Pid}, State) ->
    State2 = do_checkin(Pid, State),
    {noreply, State2};

handle_cast({set_maxconn, MaxConn}, State) ->
    {noreply, State#state{max_connections=MaxConn}};

handle_cast({set_timeout, NewTimeout}, State) ->
    %% Cap at 2 seconds
    Capped = min(NewTimeout, ?DEFAULT_KEEPALIVE_TIMEOUT),
    {noreply, State#state{keepalive_timeout=Capped}};

handle_cast({prewarm, Host, Port}, #state{prewarm_count=Count}=State) ->
    State2 = do_prewarm(Host, Port, Count, State),
    {noreply, State2};

handle_cast({prewarm, Host, Port, Count}, State) ->
    State2 = do_prewarm(Host, Port, Count, State),
    {noreply, State2};

handle_cast({prewarm_checkin, Pid, Key}, State) ->
    %% Add a prewarmed connection to the pool
    #state{available=Available, pid_monitors=PidMonitors} = State,
    %% Set owner to pool
    hackney_conn:set_owner_async(Pid, self()),
    %% Monitor the connection
    MonRef = erlang:monitor(process, Pid),
    PidMonitors2 = maps:put(Pid, MonRef, PidMonitors),
    %% Add to available
    Available2 = maps:update_with(Key, fun(Pids) -> [Pid | Pids] end, [Pid], Available),
    {noreply, State#state{available=Available2, pid_monitors=PidMonitors2}};

handle_cast({register_h2, Key, Pid}, State) ->
    %% Register an HTTP/2 connection for sharing
    #state{h2_connections = H2Conns, pid_monitors = PidMonitors} = State,
    %% Monitor the connection if not already monitored
    PidMonitors2 = case maps:is_key(Pid, PidMonitors) of
        true -> PidMonitors;
        false ->
            MonRef = erlang:monitor(process, Pid),
            maps:put(Pid, MonRef, PidMonitors)
    end,
    %% Store HTTP/2 connection
    H2Conns2 = maps:put(Key, Pid, H2Conns),
    {noreply, State#state{h2_connections = H2Conns2, pid_monitors = PidMonitors2}};

handle_cast({unregister_h2, Pid}, State) ->
    %% Remove an HTTP/2 connection from the pool
    State2 = do_unregister_h2(Pid, State),
    {noreply, State2};

handle_cast({register_h3, Key, Pid}, State) ->
    %% Register an HTTP/3 connection for sharing
    #state{h3_connections = H3Conns, pid_monitors = PidMonitors} = State,
    %% Monitor the connection if not already monitored
    PidMonitors2 = case maps:is_key(Pid, PidMonitors) of
        true -> PidMonitors;
        false ->
            MonRef = erlang:monitor(process, Pid),
            maps:put(Pid, MonRef, PidMonitors)
    end,
    %% Store HTTP/3 connection
    H3Conns2 = maps:put(Key, Pid, H3Conns),
    {noreply, State#state{h3_connections = H3Conns2, pid_monitors = PidMonitors2}};

handle_cast({unregister_h3, Pid}, State) ->
    %% Remove an HTTP/3 connection from the pool
    State2 = do_unregister_h3(Pid, State),
    {noreply, State2};

handle_cast({store_h3_session, Key, Ticket}, #state{h3_sessions = Sessions} = State) ->
    {noreply, State#state{h3_sessions = maps:put(Key, Ticket, Sessions)}};

handle_cast({delete_h3_session, Key}, #state{h3_sessions = Sessions} = State) ->
    {noreply, State#state{h3_sessions = maps:remove(Key, Sessions)}};

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info({'DOWN', _MonRef, process, Pid, Reason}, State) ->
    %% Connection process died, remove from available, in_use, h2_connections, and h3_connections
    #state{name=PoolName, available=Available, in_use=InUse, pid_monitors=PidMonitors,
           h2_connections=H2Conns, h3_connections=H3Conns} = State,
    ?report_trace("pool: connection DOWN", [{pool, PoolName}, {pid, Pid}, {reason, Reason},
        {was_in_use, maps:is_key(Pid, InUse)}]),

    %% Remove from available pools
    Available2 = maps:fold(
        fun(Key, Pids, Acc) ->
            case lists:delete(Pid, Pids) of
                [] -> maps:remove(Key, Acc);
                Pids2 -> maps:put(Key, Pids2, Acc)
            end
        end,
        Available,
        Available
    ),

    %% Remove from in_use and release load_regulation slot if it was checked out
    InUse2 = case maps:take(Pid, InUse) of
        {Key, NewInUse} ->
            %% Connection died while in use - release the load_regulation slot
            {Host, Port, _Transport} = Key,
            hackney_load_regulation:release(Host, Port),
            NewInUse;
        error ->
            InUse
    end,

    %% Remove from HTTP/2 connections if present
    H2Conns2 = maps:fold(
        fun(Key, ConnPid, Acc) ->
            case ConnPid of
                Pid -> maps:remove(Key, Acc);
                _ -> Acc
            end
        end,
        H2Conns,
        H2Conns
    ),

    %% Remove from HTTP/3 connections if present
    H3Conns2 = maps:fold(
        fun(Key, ConnPid, Acc) when is_pid(ConnPid) ->
            case ConnPid of
                Pid -> maps:remove(Key, Acc);
                _ -> Acc
            end;
           (_Key, _Val, Acc) -> Acc
        end,
        H3Conns,
        H3Conns
    ),

    PidMonitors2 = maps:remove(Pid, PidMonitors),
    {noreply, State#state{available=Available2, in_use=InUse2, pid_monitors=PidMonitors2,
                          h2_connections=H2Conns2, h3_connections=H3Conns2}};

handle_info(_Info, State) ->
    {noreply, State}.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

terminate(_Reason, #state{available=Available,
                         h2_connections=H2Conns, h3_connections=H3Conns,
                         pid_monitors=PidMonitors}) ->
    %% Stop all available connections
    maps:foreach(
        fun(_Key, Pids) ->
            lists:foreach(fun(Pid) ->
                stop_conn(Pid)
            end, Pids)
        end,
        Available
    ),

    %% Stop all HTTP/2 connections
    maps:foreach(
        fun(_Key, Pid) ->
            stop_conn(Pid)
        end,
        H2Conns
    ),

    %% Stop all HTTP/3 connections
    maps:foreach(
        fun(_Key, Pid) ->
            stop_conn(Pid)
        end,
        H3Conns
    ),

    %% Demonitor all
    maps:foreach(fun(_Pid, MonRef) -> erlang:demonitor(MonRef, [flush]) end, PidMonitors),
    ok.

%%====================================================================
%% Internal functions
%%====================================================================

connection_key(Host0, Port, Transport) ->
    Host = string:lowercase(Host0),
    {Host, Port, Transport}.

%% @private Stop a connection, tolerating an already-dead process.
stop_conn(Pid) ->
    try hackney_conn:stop(Pid) catch _:_ -> ok end.

find_available(Key, Available) ->
    case maps:find(Key, Available) of
        {ok, [Pid | Rest]} ->
            Available2 = case Rest of
                [] -> maps:remove(Key, Available);
                _ -> maps:put(Key, Rest, Available)
            end,
            %% Verify connection is still alive and usable
            case is_process_alive(Pid) of
                true ->
                    %% is_ready checks both state and socket health in one call.
                    %% The connection can die between is_process_alive/1 above
                    %% and these gen_statem calls (flaky network); the resulting
                    %% noproc exit must not crash the pool, so skip and move on.
                    try hackney_conn:is_ready(Pid) of
                        {ok, connected} -> {ok, Pid, Available2};
                        {ok, closed} ->
                            %% Connection closed, try reconnect
                            try hackney_conn:connect(Pid) of
                                ok -> {ok, Pid, Available2};
                                _ -> find_available(Key, Available2)
                            catch
                                _:_ -> find_available(Key, Available2)
                            end;
                        _ -> find_available(Key, Available2)
                    catch
                        _:_ -> find_available(Key, Available2)
                    end;
                false ->
                    find_available(Key, Available2)
            end;
        {ok, []} ->
            none;
        error ->
            none
    end.

start_connection(Key, Owner, Opts, State) ->
    {Host, Port, Transport} = Key,
    ConnectTimeout = proplists:get_value(connect_timeout, Opts, 8000),
    RecvTimeout = proplists:get_value(recv_timeout, Opts, infinity),
    IdleTimeout = State#state.keepalive_timeout,
    SslOpts = proplists:get_value(ssl_options, Opts, []),
    ConnectOpts = proplists:get_value(connect_options, Opts, []),

    ConnOpts = #{
        host => Host,
        port => Port,
        transport => Transport,
        connect_timeout => ConnectTimeout,
        recv_timeout => RecvTimeout,
        idle_timeout => IdleTimeout,
        ssl_options => SslOpts,
        connect_options => ConnectOpts,
        pool_pid => self(),
        owner => Owner
    },

    case hackney_conn_sup:start_conn(ConnOpts) of
        {ok, Pid} ->
            %% Connect the connection
            case hackney_conn:connect(Pid) of
                ok ->
                    %% Monitor the process
                    MonRef = erlang:monitor(process, Pid),
                    PidMonitors = maps:put(Pid, MonRef, State#state.pid_monitors),
                    {ok, Pid, State#state{pid_monitors=PidMonitors}};
                {error, Reason} ->
                    stop_conn(Pid),
                    {error, Reason}
            end;
        {error, Reason} ->
            {error, Reason}
    end.

%% @private Process a checkin - return connection to pool
%% Only TCP connections are stored. SSL upgraded connections are closed.
%% Always releases the load_regulation slot since connection is no longer in use.
do_checkin(Pid, State) ->
    #state{in_use=InUse, available=Available, pid_monitors=PidMonitors} = State,

    %% Get the key from in_use and remove
    case maps:take(Pid, InUse) of
        {Key, InUse2} ->
            %% Release load_regulation slot - connection is no longer in active use
            {Host, Port, _Transport} = Key,
            hackney_load_regulation:release(Host, Port),

            %% Check if connection is still alive
            case is_process_alive(Pid) of
                true ->
                    %% Check if this connection should not be reused:
                    %% - SSL upgraded connections (security requirement)
                    %% - Proxy tunnel connections (SOCKS5, HTTP CONNECT - issue #283)
                    ShouldClose = (try hackney_conn:is_upgraded_ssl(Pid) catch _:_ -> false end) =:= true orelse
                                  (try hackney_conn:is_no_reuse(Pid) catch _:_ -> false end) =:= true,
                    case ShouldClose of
                        true ->
                            %% Connection should not be reused - close it
                            stop_conn(Pid),
                            %% Remove monitor if exists
                            PidMonitors2 = case maps:take(Pid, PidMonitors) of
                                {MonRef, PM} ->
                                    erlang:demonitor(MonRef, [flush]),
                                    PM;
                                error -> PidMonitors
                            end,
                            State2 = State#state{in_use=InUse2, pid_monitors=PidMonitors2},
                            %% Still trigger prewarm for TCP connections (for future SSL upgrades)
                            TcpKey = {Host, Port, hackney_tcp},
                            maybe_maintain_prewarm(TcpKey, activate_host(TcpKey, State2));
                        _ ->
                            %% Regular TCP connection - store in pool
                            %% Set owner to pool so connection doesn't die if previous requester crashes.
                            %% Use async version to avoid deadlock when called from checkin_sync
                            hackney_conn:set_owner_async(Pid, self()),
                            Available2 = maps:update_with(Key, fun(Pids) -> [Pid | Pids] end, [Pid], Available),

                            %% Ensure we're monitoring this pid
                            PidMonitors2 = case maps:is_key(Pid, PidMonitors) of
                                true -> PidMonitors;
                                false ->
                                    MonRef = erlang:monitor(process, Pid),
                                    maps:put(Pid, MonRef, PidMonitors)
                            end,

                            State2 = State#state{in_use=InUse2, available=Available2, pid_monitors=PidMonitors2},
                            %% Maintain prewarm for activated hosts
                            maybe_maintain_prewarm(Key, State2)
                    end;
                false ->
                    State#state{in_use=InUse2}
            end;
        error ->
            State
    end.

%% @private Process a checkin with known "should close" flag - avoids calling back to connection
%% Used for sync checkin to prevent deadlock when connection calls pool.
%% ShouldClose is true if connection was SSL upgraded or is a proxy tunnel (no_reuse).
do_checkin_with_close_flag(Pid, ShouldClose, State) ->
    #state{in_use=InUse, available=Available, pid_monitors=PidMonitors} = State,

    %% Get the key from in_use and remove
    case maps:take(Pid, InUse) of
        {Key, InUse2} ->
            %% Release load_regulation slot - connection is no longer in active use
            {Host, Port, _Transport} = Key,
            hackney_load_regulation:release(Host, Port),

            %% Check if connection is still alive
            case is_process_alive(Pid) of
                true ->
                    %% Use the provided flag instead of querying connection
                    case ShouldClose of
                        true ->
                            %% Connection should not be reused - close it
                            %% Use cast to avoid deadlock (connection is waiting for our reply)
                            gen_statem:cast(Pid, stop),
                            %% Remove monitor if exists
                            PidMonitors2 = case maps:take(Pid, PidMonitors) of
                                {MonRef, PM} ->
                                    erlang:demonitor(MonRef, [flush]),
                                    PM;
                                error -> PidMonitors
                            end,
                            State2 = State#state{in_use=InUse2, pid_monitors=PidMonitors2},
                            %% Still trigger prewarm for TCP connections (for future SSL upgrades)
                            TcpKey = {Host, Port, hackney_tcp},
                            maybe_maintain_prewarm(TcpKey, activate_host(TcpKey, State2));
                        _ ->
                            %% Regular connection - store in pool
                            %% Set owner to pool so connection doesn't die if previous requester crashes.
                            hackney_conn:set_owner_async(Pid, self()),
                            Available2 = maps:update_with(Key, fun(Pids) -> [Pid | Pids] end, [Pid], Available),

                            %% Ensure we're monitoring this pid
                            PidMonitors2 = case maps:is_key(Pid, PidMonitors) of
                                true -> PidMonitors;
                                false ->
                                    MonRef = erlang:monitor(process, Pid),
                                    maps:put(Pid, MonRef, PidMonitors)
                            end,

                            State2 = State#state{in_use=InUse2, available=Available2, pid_monitors=PidMonitors2},
                            %% Maintain prewarm for activated hosts
                            maybe_maintain_prewarm(Key, State2)
                    end;
                false ->
                    State#state{in_use=InUse2}
            end;
        error ->
            State
    end.

%% @private Check that a pooled HTTP/2 conn is alive and in `connected` state.
%% Short timeout so a stuck conn doesn't wedge the pool; any failure → unusable.
h2_conn_usable(Pid) ->
    case erlang:is_process_alive(Pid) of
        false -> false;
        true ->
            try hackney_conn:get_state(Pid) of
                {ok, connected} -> true;
                _ -> false
            catch
                _:_ -> false
            end
    end.

%% @private Remove an HTTP/2 connection from the pool
do_unregister_h2(Pid, State) ->
    #state{h2_connections = H2Conns, pid_monitors = PidMonitors} = State,
    %% Find and remove the connection
    H2Conns2 = maps:fold(
        fun(Key, ConnPid, Acc) ->
            case ConnPid of
                Pid -> maps:remove(Key, Acc);
                _ -> Acc
            end
        end,
        H2Conns,
        H2Conns
    ),
    %% Demonitor if no longer tracked
    PidMonitors2 = case maps:take(Pid, PidMonitors) of
        {MonRef, PM} ->
            erlang:demonitor(MonRef, [flush]),
            PM;
        error -> PidMonitors
    end,
    State#state{h2_connections = H2Conns2, pid_monitors = PidMonitors2}.

%% @private Remove an HTTP/3 connection from the pool
do_unregister_h3(Pid, State) ->
    #state{h3_connections = H3Conns, pid_monitors = PidMonitors} = State,
    %% Find and remove the connection
    H3Conns2 = maps:fold(
        fun(Key, ConnPid, Acc) when is_pid(ConnPid) ->
            case ConnPid of
                Pid -> maps:remove(Key, Acc);
                _ -> Acc
            end;
           (_Key, _Val, Acc) -> Acc
        end,
        H3Conns,
        H3Conns
    ),
    %% Demonitor if no longer tracked
    PidMonitors2 = case maps:take(Pid, PidMonitors) of
        {MonRef, PM} ->
            erlang:demonitor(MonRef, [flush]),
            PM;
        error -> PidMonitors
    end,
    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) ->
    #state{name=PoolName, available=Available, keepalive_timeout=IdleTimeout} = State,
    Key = connection_key(Host, Port, hackney_tcp),
    CurrentCount = case maps:find(Key, Available) of
        {ok, Pids} -> length(Pids);
        error -> 0
    end,
    Needed = max(0, Count - CurrentCount),
    ?report_trace("pool: prewarm", [{pool, PoolName}, {host, Host}, {port, Port},
                                    {current, CurrentCount}, {needed, Needed}]),
    case Needed > 0 of
        true ->
            %% Mark host as activated
            State2 = activate_host(Key, State),
            %% Create connections in background to not block the pool
            PoolPid = self(),
            spawn(fun() ->
                prewarm_connections(PoolPid, Host, Port, Needed, IdleTimeout)
            end),
            State2;
        false ->
            State
    end.

%% @private Mark a host as activated for prewarm
activate_host(Key, #state{activated_hosts=Activated}=State) ->
    State#state{activated_hosts=sets:add_element(Key, Activated)}.

%% @private Check if a host has been activated
is_host_activated(Key, #state{activated_hosts=Activated}) ->
    sets:is_element(Key, Activated).

%% @private Maintain prewarm count for activated hosts
maybe_maintain_prewarm(Key, State) ->
    case is_host_activated(Key, State) of
        true ->
            #state{prewarm_count=PrewarmCount, available=Available} = State,
            CurrentCount = case maps:find(Key, Available) of
                {ok, Pids} -> length(Pids);
                error -> 0
            end,
            case CurrentCount < PrewarmCount of
                true ->
                    {Host, Port, _Transport} = Key,
                    do_prewarm(Host, Port, PrewarmCount, State);
                false ->
                    State
            end;
        false ->
            State
    end.

%% @private Create prewarm connections (runs in separate process)
prewarm_connections(_PoolPid, _Host, _Port, 0, _IdleTimeout) ->
    ok;
prewarm_connections(PoolPid, Host, Port, Count, IdleTimeout) ->
    ConnOpts = #{
        host => Host,
        port => Port,
        transport => hackney_tcp,
        connect_timeout => 5000,
        recv_timeout => 5000,
        idle_timeout => IdleTimeout,
        ssl_options => [],
        connect_options => [],
        pool_pid => PoolPid,
        owner => PoolPid
    },
    case hackney_conn_sup:start_conn(ConnOpts) of
        {ok, Pid} ->
            case hackney_conn:connect(Pid) of
                ok ->
                    %% Checkin the new connection to the pool
                    gen_server:cast(PoolPid, {prewarm_checkin, Pid, {Host, Port, hackney_tcp}});
                {error, _Reason} ->
                    stop_conn(Pid)
            end;
        {error, _Reason} ->
            ok
    end,
    prewarm_connections(PoolPid, Host, Port, Count - 1, IdleTimeout).

handle_stats(State) ->
    #state{name=PoolName, max_connections=Max, available=Available,
           in_use=InUse} = State,
    AvailCount = maps:fold(fun(_, Pids, Acc) -> Acc + length(Pids) end, 0, Available),
    [{name, PoolName},
     {max, Max},
     {in_use_count, maps:size(InUse)},
     {free_count, AvailCount},
     {queue_count, 0}].  % No more queue