Fix per-host connection permit leak on the connect path

client/src/main/java/org/asynchttpclient/netty/NettyResponseFuture.java

@@ -117,8 +117,13 @@ public final class NettyResponseFuture<V> implements ListenableFuture<V> {
     // volatile where we don't need CAS ops
     private volatile long touch = unpreciseMillisTime();
     private volatile ChannelState channelState = ChannelState.NEW;
+    // Written and read from the event loop, the caller thread (attachChannel on a pooled hit, cancel()) and
+    // the HashedWheelTimer thread (terminateAndExit, and TimeoutTimerTask.expire which closes this channel
+    // on a request timeout). Without volatile those readers can miss the write and skip the close, which
+    // since the connect path publishes the channel before the TLS handshake would strand a stuck
+    // connection, and its permit, until handshakeTimeout (issue #2189).
+    private volatile Channel channel;
     // state mutated only inside the event loop
-    private Channel channel;
     private boolean keepAlive = true;
     private Request targetRequest;
     private Request currentRequest;

client/src/main/java/org/asynchttpclient/netty/channel/Http2ConnectionState.java

@@ -77,9 +77,10 @@ private static final class PendingOpener {
     private final AtomicBoolean redundant = new AtomicBoolean(false);
     private volatile Object partitionKey;
     // Releases the per-host permit this connection holds in ROUND_ROBIN mode (installed by
-    // NettyConnectListener, absent in DEFAULT mode). The GOAWAY handler and the channel closeFuture race
-    // to fire it; getAndSet(null) makes it run at most once, since a double release would push the
-    // semaphore above maxConnectionsPerHost.
+    // NettyConnectListener, absent in DEFAULT mode). Fired by the GOAWAY handler so the permit is freed at
+    // drain start rather than at close; the channel closeFuture releases the same permit directly, and both
+    // funnel through a single getAndSet, since a double release would push the semaphore above
+    // maxConnectionsPerHost.
     private final AtomicReference<Runnable> permitRelease = new AtomicReference<>();
 
     public boolean tryAcquireStream() {
@@ -281,8 +282,8 @@ public void setPermitRelease(Runnable release) {
 
     /**
      * Runs the installed permit-release action the first time it is called; later calls do nothing, as do
-     * calls when no action was installed (DEFAULT mode). Fired at drain start (GOAWAY) and on channel
-     * close, whichever comes first.
+     * calls when no action was installed (DEFAULT mode). Fired at drain start (GOAWAY); a channel that
+     * closes without draining has its permit released by the closeFuture listener instead.
      */
     public void releasePermitOnce() {
         Runnable release = permitRelease.getAndSet(null);

client/src/main/java/org/asynchttpclient/netty/channel/NettyConnectListener.java

@@ -34,6 +34,7 @@
 
 import java.net.ConnectException;
 import java.net.InetSocketAddress;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Non Blocking connect.
@@ -75,22 +76,34 @@ private void writeRequest(Channel channel) {
             LOGGER.debug("Using new Channel '{}' for '{}' to '{}'", channel, httpRequest.method(), httpRequest.uri());
         }
 
+        // Publishing the future on the channel is deliberately deferred to here, unlike attachChannel which
+        // runs before the handshake: it arms AsyncHttpClientHandler.channelInactive, whose
+        // handleUnexpectedClosedChannel would race NettyConnectListener.onFailure to retry the same connect
+        // failure twice if a mid-handshake close could reach it.
         Channels.setAttribute(channel, future);
 
         channelManager.registerOpenChannel(channel);
-        future.attachChannel(channel, false);
         requestSender.writeRequest(future, channel);
     }
 
     public void onSuccess(Channel channel, InetSocketAddress remoteAddress) {
-        // Take the semaphore lock from the future. For HTTP/1.1, we'll transfer it to channel.closeFuture().
-        // For HTTP/2, we release it immediately after ALPN negotiation since the connection is multiplexed.
-        final Object partitionKeyLock = (connectionSemaphore != null) ? future.takePartitionKeyLock() : null;
+        // The permit must be owned by the channel from the moment it leaves the future: takePartitionKeyLock()
+        // is a getAndSet(null), so abort() can no longer return it, and everything below can abandon this
+        // channel without being able to reach the token -- a failed or timed-out TLS handshake, a crashing
+        // AsyncHandler callback, a connect retry (issue #2189). The HTTP/2 branches release it earlier and
+        // leave this listener a no-op. The listener captures no enclosing instance, so it does not pin this
+        // request's object graph for the lifetime of a long-lived (HTTP/2) connection.
+        final ConnectionSemaphore semaphore = connectionSemaphore;
+        final Object partitionKeyLock = semaphore != null ? future.takePartitionKeyLock() : null;
+        final AtomicReference<Object> permit = new AtomicReference<>(partitionKeyLock);
+        if (partitionKeyLock != null) {
+            channel.closeFuture().addListener(f -> releasePermitOnce(semaphore, permit));
+        }
 
         Channels.setActiveToken(channel);
 
         if (futureIsAlreadyCompleted(channel)) {
-            releaseSemaphoreImmediately(partitionKeyLock);
+            releasePermitOnce(semaphore, permit);
             return;
         }
 
@@ -102,10 +115,15 @@ public void onSuccess(Channel channel, InetSocketAddress remoteAddress) {
             // futureIsAlreadyCompleted check above can pass while the holder is already null.
             // Drop this connection rather than NPE-ing on setResolvedRemoteAddress below.
             Channels.silentlyCloseChannel(channel);
-            releaseSemaphoreImmediately(partitionKeyLock);
+            releasePermitOnce(semaphore, permit);
             return;
         }
 
+        // Publish the channel before the (possibly long) TLS handshake. Until this runs future.channel() is
+        // null, and NettyRequestSender.abort only closes a non-null channel -- so a request timeout firing
+        // mid-handshake could not close the socket, stranding it until handshakeTimeout (issue #2189).
+        future.attachChannel(channel, false);
+
         Request request = future.getTargetRequest();
         Uri uri = request.getUri();
         // don't set a null resolved address - if the remoteAddress is null we keep
@@ -153,7 +171,6 @@ protected void onSuccess(Channel value) {
                         return;
                     }
                     // After SSL handshake to proxy, continue with normal proxy request
-                    attachSemaphoreToChannelClose(channel, partitionKeyLock);
                     writeRequest(channel);
                 }
 
@@ -215,9 +232,7 @@ protected void onSuccess(Channel value) {
                     }
                     if (http2Negotiated && !uri.isWebSocket()) {
                         channelManager.upgradePipelineToHttp2(channel.pipeline());
-                        registerHttp2AndManageSemaphore(channel, partitionKeyLock);
-                    } else {
-                        attachSemaphoreToChannelClose(channel, partitionKeyLock);
+                        registerHttp2AndManageSemaphore(channel, semaphore, permit);
                     }
                     writeRequest(channel);
                 }
@@ -240,32 +255,22 @@ protected void onFailure(Throwable cause) {
             // excluded for the same RFC 8441 reason as the TLS path above — it stays on HTTP/1.1.
             if (!uri.isSecured() && channelManager.isHttp2CleartextEnabled() && !uri.isWebSocket()) {
                 channelManager.upgradePipelineToHttp2(channel.pipeline());
-                registerHttp2AndManageSemaphore(channel, partitionKeyLock);
-            } else {
-                attachSemaphoreToChannelClose(channel, partitionKeyLock);
+                registerHttp2AndManageSemaphore(channel, semaphore, permit);
             }
             writeRequest(channel);
         }
     }
 
     /**
-     * Attaches the semaphore lock to the channel's close future (HTTP/1.1 behavior).
-     * The semaphore slot is released when the connection closes.
-     */
-    private void attachSemaphoreToChannelClose(Channel channel, Object partitionKeyLock) {
-        if (connectionSemaphore != null && partitionKeyLock != null) {
-            channel.closeFuture().addListener(f -> connectionSemaphore.releaseChannelLock(partitionKeyLock));
-        }
-    }
-
-    /**
-     * Releases the semaphore lock immediately (HTTP/2 behavior).
-     * HTTP/2 connections are multiplexed, so the semaphore should not be held
-     * for the lifetime of the connection.
+     * Returns the per-host permit this connection holds, at most once. The channel closeFuture, the
+     * early-exit paths and the HTTP/2 immediate release all race to call this; a double release would
+     * push the semaphore above {@code maxConnectionsPerHost}, and an unmatched one can prematurely
+     * prune a live entry from the per-host map.
      */
-    private void releaseSemaphoreImmediately(Object partitionKeyLock) {
-        if (connectionSemaphore != null && partitionKeyLock != null) {
-            connectionSemaphore.releaseChannelLock(partitionKeyLock);
+    private static void releasePermitOnce(ConnectionSemaphore semaphore, AtomicReference<Object> permit) {
+        Object partitionKeyLock = permit.getAndSet(null);
+        if (partitionKeyLock != null && semaphore != null) {
+            semaphore.releaseChannelLock(partitionKeyLock);
         }
     }
 
@@ -281,27 +286,22 @@ private void releaseSemaphoreImmediately(Object partitionKeyLock) {
      * a further request fails to acquire a permit and multiplexes onto a sibling-IP connection via
      * {@link ChannelManager#pollHttp2SiblingConnection(Object)} rather than opening another (issue #2214).
      */
-    private void registerHttp2AndManageSemaphore(Channel channel, Object partitionKeyLock) {
+    private void registerHttp2AndManageSemaphore(Channel channel, ConnectionSemaphore semaphore, AtomicReference<Object> permit) {
         // Register under the future's partition key so the H2 connection is found by the same key the
         // pool is polled with, including the IP-aware key used by LoadBalance.ROUND_ROBIN. Read the key
         // once: it is volatile (repinned on IP failover) and both uses below must agree.
         Object partitionKey = future.getPartitionKey();
         channelManager.registerHttp2Connection(partitionKey, channel);
         if (partitionKey instanceof RoundRobinPartitionKey) {
             // The permit must be freed as soon as the connection stops serving new requests: on close, or
-            // at drain start after GOAWAY (issue #2214). Both paths go through releasePermitOnce(), which
-            // guarantees a single release; a double release would push the semaphore above the cap. The
-            // state attribute is always present after upgradePipelineToHttp2, the fallback is just
-            // belt and braces.
+            // at drain start after GOAWAY (issue #2214). The closeFuture listener installed in onSuccess
+            // covers close; this adds the drain hook. Both funnel through releasePermitOnce.
             Http2ConnectionState state = channel.attr(Http2ConnectionState.HTTP2_STATE_KEY).get();
-            if (state != null && connectionSemaphore != null && partitionKeyLock != null) {
-                state.setPermitRelease(() -> connectionSemaphore.releaseChannelLock(partitionKeyLock));
-                channel.closeFuture().addListener(f -> state.releasePermitOnce());
-            } else {
-                attachSemaphoreToChannelClose(channel, partitionKeyLock);
+            if (state != null) {
+                state.setPermitRelease(() -> releasePermitOnce(semaphore, permit));
             }
         } else {
-            releaseSemaphoreImmediately(partitionKeyLock);
+            releasePermitOnce(semaphore, permit);
         }
     }
 

client/src/main/java/org/asynchttpclient/netty/handler/WebSocketHandler.java

@@ -107,6 +107,16 @@ private void abort(Channel channel, NettyResponseFuture<?> future, WebSocketUpgr
     public void handleRead(Channel channel, NettyResponseFuture<?> future, Object e) throws Exception {
 
         if (e instanceof HttpResponse) {
+            // Unlike HttpHandler, this check cannot guard the whole method: a successful upgrade completes
+            // the future (see upgrade() below) and the channel then keeps serving frames, so isDone() is
+            // true for all normal WebSocket traffic. It belongs on the upgrade response alone, where a 101
+            // arriving just as a request timeout aborted the future would otherwise still run onOpen and
+            // deliver the WebSocket lifecycle after onThrowable had already fired.
+            if (future.isDone()) {
+                channelManager.closeChannel(channel);
+                return;
+            }
+
             HttpResponse response = (HttpResponse) e;
             if (logger.isDebugEnabled()) {
                 HttpRequest httpRequest = future.getNettyRequest().getHttpRequest();

client/src/main/java/org/asynchttpclient/netty/request/NettyRequestSender.java

@@ -1071,18 +1071,21 @@ private static void scheduleReadTimeout(NettyResponseFuture<?> nettyResponseFutu
     }
 
     public void abort(Channel channel, NettyResponseFuture<?> future, Throwable t) {
-        if (channel != null) {
-            if (channel.isActive()) {
-                channelManager.closeChannel(channel);
-            }
-        }
-
+        // Complete the future before closing, so the caller's cause is the one the user sees. Closing first
+        // can fail an in-flight TLS handshake, and that failure races back through
+        // NettyConnectListener.onFailure to abort the same future with a ConnectException instead -- which a
+        // request timeout on the connect path can now hit, since the channel is published before the
+        // handshake (issue #2189). The close still uses the channel passed in, which abort() does not clear.
         if (!future.isDone()) {
             future.setChannelState(ChannelState.CLOSED);
             LOGGER.debug("Aborting Future {}\n", future);
             LOGGER.debug(t.getMessage(), t);
             future.abort(t);
         }
+
+        if (channel != null && channel.isActive()) {
+            channelManager.closeChannel(channel);
+        }
     }
 
     public void handleUnexpectedClosedChannel(Channel channel, NettyResponseFuture<?> future) {

client/src/test/java/org/asynchttpclient/netty/channel/ChannelManagerHttp2DrainPermitTest.java

@@ -27,6 +27,7 @@
 
 import java.net.InetAddress;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.asynchttpclient.Dsl.config;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
@@ -81,9 +82,19 @@ private EmbeddedChannel registerRrConnectionHoldingPermit(ConnectionSemaphore se
         Http2ConnectionState state = new Http2ConnectionState();
         channel.attr(Http2ConnectionState.HTTP2_STATE_KEY).set(state);
         channelManager.registerHttp2Connection(registryKey, channel);
-        // Mirror NettyConnectListener.registerHttp2AndManageSemaphore's round-robin branch.
-        state.setPermitRelease(() -> semaphore.releaseChannelLock(baseKey));
-        channel.closeFuture().addListener(f -> state.releasePermitOnce());
+        // Stand in for NettyConnectListener's round-robin wiring - the drain hook from
+        // registerHttp2AndManageSemaphore plus the closeFuture release from onSuccess, both funnelling
+        // through one getAndSet. This fixture only pins ChannelManager's GOAWAY drain contract; that
+        // NettyConnectListener actually installs this wiring is covered end-to-end by BasicHttp2Test.
+        AtomicReference<Object> permit = new AtomicReference<>(baseKey);
+        Runnable release = () -> {
+            Object key = permit.getAndSet(null);
+            if (key != null) {
+                semaphore.releaseChannelLock(key);
+            }
+        };
+        state.setPermitRelease(release);
+        channel.closeFuture().addListener(f -> release.run());
         return channel;
     }