Reuse the per-request partition key instead of recomputing it

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

@@ -136,6 +136,11 @@ public final class NettyResponseFuture<V> implements ListenableFuture<V> {
     private volatile List<InetSocketAddress> roundRobinAddresses;
     private volatile Uri roundRobinBaseUri;
     private volatile ScramContext scramContext;
+    // Base (host/scheme/port) partition key, computed eagerly at construction and recomputed by
+    // setTargetRequest (its only mutator: connectionPoolPartitioning/proxyServer are final and targetRequest
+    // is its only other input). Volatile: setTargetRequest runs on the redirect path while reads happen on
+    // other threads.
+    private volatile Object basePartitionKeyCache;
 
     public NettyResponseFuture(Request originalRequest,
                                AsyncHandler<V> asyncHandler,
@@ -152,6 +157,7 @@ public NettyResponseFuture(Request originalRequest,
         this.connectionSemaphore = connectionSemaphore;
         this.proxyServer = proxyServer;
         this.maxRetry = maxRetry;
+        basePartitionKeyCache = computeBasePartitionKey();
     }
 
     private void releasePartitionKeyLock() {
@@ -370,6 +376,9 @@ public Request getTargetRequest() {
 
     public void setTargetRequest(Request targetRequest) {
         this.targetRequest = targetRequest;
+        // Recompute the base partition key eagerly: a redirect/retry may target a different
+        // host/scheme/port, which changes the key.
+        basePartitionKeyCache = computeBasePartitionKey();
     }
 
     public Request getCurrentRequest() {
@@ -546,14 +555,32 @@ public Object getPartitionKey() {
     }
 
     /**
-     * The per-host partition key, ignoring any round-robin IP-aware override. Used for the connection
-     * semaphore so {@code maxConnectionsPerHost} stays per host (not per IP): the permit is taken
-     * before the target IP is known and the connector may fail over to a different IP than the one
-     * initially selected.
+     * The per-host base partition key, ignoring any round-robin IP-aware override. This is the key used to
+     * acquire the connection semaphore (so {@code maxConnectionsPerHost} stays per host, not per IP: the
+     * permit is taken before the target IP is known and the connector may fail over to a different IP than
+     * the one initially selected), to offer the channel back to the pool, and to register/poll the HTTP/2
+     * connection registry; in round-robin mode it is also the base that the per-IP override
+     * ({@link RoundRobinPartitionKey}) wraps.
+     * <p>
+     * Note: the pool/H2 <em>poll</em> paths ({@link org.asynchttpclient.netty.request.NettyRequestSender}
+     * pollPooledChannel/waitForHttp2Connection) intentionally derive their key from the live request rather
+     * than this value, because a filter replay can reuse a future without updating its {@code targetRequest}.
+     * <p>
+     * Computed eagerly at construction and recomputed by {@link #setTargetRequest(Request)}; this accessor
+     * is a plain read of the memoized value.
      */
     public Object basePartitionKey() {
-        return connectionPoolPartitioning.getPartitionKey(targetRequest.getUri(), targetRequest.getVirtualHost(),
-                proxyServer);
+        return basePartitionKeyCache;
+    }
+
+    // Depends only on targetRequest (host/scheme/port + virtualHost) and the final proxyServer, so it is
+    // recomputed only when setTargetRequest changes the target. Tolerates null inputs: some unit tests
+    // construct a future with no request/partitioning and never consult the key; production always has both.
+    private Object computeBasePartitionKey() {
+        return targetRequest != null && connectionPoolPartitioning != null
+                ? connectionPoolPartitioning.getPartitionKey(targetRequest.getUri(), targetRequest.getVirtualHost(),
+                        proxyServer)
+                : null;
     }
 
     /**

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

@@ -442,12 +442,24 @@ public Channel pollHttp2Connection(Object partitionKey) {
     /**
      * Polls for an HTTP/2 connection by URI/virtualHost/proxy, using the same partition key logic
      * as the regular pool. Returns the connection without removing it from the registry.
+     *
+     * @deprecated no longer used internally. Compute the partition key at the call site — from the request
+     * being dispatched, so it stays correct on the filter-replay path — and call
+     * {@link #pollHttp2Connection(Object)}. Kept for binary compatibility; slated for removal in the next
+     * major release.
      */
+    @Deprecated
     public Channel pollHttp2(Uri uri, String virtualHost, ProxyServer proxy, ChannelPoolPartitioning connectionPoolPartitioning) {
         Object partitionKey = connectionPoolPartitioning.getPartitionKey(uri, virtualHost, proxy);
         return pollHttp2Connection(partitionKey);
     }
 
+    /**
+     * @deprecated no longer used internally. Compute the partition key at the call site — from the request
+     * being dispatched, so it stays correct on the filter-replay path — and call {@link #poll(Object)}.
+     * Kept for binary compatibility; slated for removal in the next major release.
+     */
+    @Deprecated
     public Channel poll(Uri uri, String virtualHost, ProxyServer proxy, ChannelPoolPartitioning connectionPoolPartitioning) {
         Object partitionKey = connectionPoolPartitioning.getPartitionKey(uri, virtualHost, proxy);
         return channelPool.poll(partitionKey);

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

@@ -214,6 +214,10 @@ protected void onSuccess(List<InetSocketAddress> addresses) {
                 if (addresses.size() > 1) {
                     ordered = rrSelector.rotate(host, addresses);
                     InetAddress chosen = ordered.get(0).getAddress();
+                    // Derive the base key from the live request, not newFuture.basePartitionKey(): a filter
+                    // replay can reuse this future without updating its targetRequest (only redirects keep it
+                    // in sync), so the memoized base key may lag a host-rewriting replay and would mis-key the
+                    // override to the previous host.
                     Object baseKey = request.getChannelPoolPartitioning().getPartitionKey(uri, request.getVirtualHost(), proxyServer);
                     newFuture.setPartitionKeyOverride(new RoundRobinPartitionKey(baseKey, chosen));
                 } else {
@@ -1128,10 +1132,16 @@ private Channel waitForHttp2Connection(Request request, ProxyServer proxy, Netty
             return null;
         }
         String virtualHost = request.getVirtualHost();
-        // In round-robin mode, only multiplex onto the H2 connection for the IP this request is pinned to.
+        // In round-robin mode, only multiplex onto the H2 connection for the IP this request is pinned to;
+        // otherwise use the per-host base key. Derive it from the live request rather than
+        // future.basePartitionKey() so a filter replay that rewrites the host still polls the correct key
+        // (the future's targetRequest is only kept in sync on the redirect path). Computed once and reused
+        // across the poll loop below.
         Object override = future != null ? future.getPartitionKeyOverride() : null;
+        Object h2Key = override != null ? override
+                : request.getChannelPoolPartitioning().getPartitionKey(uri, virtualHost, proxy);
 
-        Channel h2Channel = pollHttp2(override, uri, virtualHost, proxy, request);
+        Channel h2Channel = channelManager.pollHttp2Connection(h2Key);
         if (h2Channel != null) {
             return h2Channel;
         }
@@ -1147,7 +1157,7 @@ private Channel waitForHttp2Connection(Request request, ProxyServer proxy, Netty
 
         long deadline = System.nanoTime() + config.getConnectTimeout().toNanos();
         while (System.nanoTime() < deadline) {
-            h2Channel = pollHttp2(override, uri, virtualHost, proxy, request);
+            h2Channel = channelManager.pollHttp2Connection(h2Key);
             if (h2Channel != null) {
                 return h2Channel;
             }
@@ -1161,14 +1171,6 @@ private Channel waitForHttp2Connection(Request request, ProxyServer proxy, Netty
         return null;
     }
 
-    // Polls the HTTP/2 registry, using the IP-aware key in round-robin mode and the regular key otherwise.
-    private Channel pollHttp2(Object override, Uri uri, String virtualHost, ProxyServer proxy, Request request) {
-        if (override != null) {
-            return channelManager.pollHttp2Connection(override);
-        }
-        return channelManager.pollHttp2(uri, virtualHost, proxy, request.getChannelPoolPartitioning());
-    }
-
     private boolean isOnEventLoop() {
         for (EventExecutor executor : channelManager.getEventLoopGroup()) {
             if (executor.inEventLoop()) {
@@ -1212,15 +1214,21 @@ private Channel pollPooledChannel(NettyResponseFuture<?> future, Request request
         // connection would send the WS handshake as a plain HTTP/2 request and the WebSocket handler would
         // receive raw frames ("Invalid message ... AdaptiveByteBuf"). Fall through to an HTTP/1.1 connection.
         // See Issue #2160.
+        // Compute the base partition key once and reuse it for both the HTTP/2 registry poll and the
+        // HTTP/1.1 pool poll, instead of recomputing (and re-allocating) it inside each channelManager call.
+        // Derive it from the live request rather than future.basePartitionKey(): on the filter-replay path
+        // (replayRequest) the future's targetRequest is not updated to the replayed request, so its memoized
+        // base key can lag a host-rewriting replay. Reading the current request's URI/virtualHost stays correct.
+        Object partitionKey = request.getChannelPoolPartitioning().getPartitionKey(uri, virtualHost, proxy);
         if (!uri.isWebSocket()) {
-            Channel h2Channel = channelManager.pollHttp2(uri, virtualHost, proxy, request.getChannelPoolPartitioning());
+            Channel h2Channel = channelManager.pollHttp2Connection(partitionKey);
             if (h2Channel != null) {
                 LOGGER.debug("Using HTTP/2 multiplexed Channel '{}' for '{}' to '{}'", h2Channel, request.getMethod(), uri);
                 return h2Channel;
             }
         }
 
-        final Channel channel = channelManager.poll(uri, virtualHost, proxy, request.getChannelPoolPartitioning());
+        final Channel channel = channelManager.poll(partitionKey);
 
         if (channel != null) {
             LOGGER.debug("Using pooled Channel '{}' for '{}' to '{}'", channel, request.getMethod(), uri);

client/src/test/java/org/asynchttpclient/netty/NettyResponseFutureTest.java

@@ -17,12 +17,18 @@
 
 import io.github.artsok.RepeatedIfExceptionsTest;
 import org.asynchttpclient.AsyncHandler;
+import org.asynchttpclient.Request;
+import org.asynchttpclient.channel.ChannelPoolPartitioning;
+import org.junit.jupiter.api.Test;
 
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 
+import static org.asynchttpclient.Dsl.get;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -90,4 +96,29 @@ public void testGetThrowsExceptionOnAbort() throws Exception {
         assertThrows(ExecutionException.class, () -> nettyResponseFuture.get(),
                 "An ExecutionException must have occurred by now as 'abort' was called before 'get'");
     }
+
+    @Test
+    public void basePartitionKeyIsMemoizedAndInvalidatedOnTargetChange() {
+        AsyncHandler<?> asyncHandler = mock(AsyncHandler.class);
+        Request reqA = get("http://hosta.example/").build();
+        ChannelPoolPartitioning partitioning = reqA.getChannelPoolPartitioning();
+        NettyResponseFuture<?> future = new NettyResponseFuture<>(reqA, asyncHandler, null, 3, partitioning, null, null);
+
+        Object k1 = future.basePartitionKey();
+        Object k2 = future.basePartitionKey();
+        // Memoized: repeat calls return the SAME instance (previously each call allocated a fresh key).
+        assertSame(k1, k2, "base partition key must be memoized (same instance on repeat calls)");
+        // ...and it equals a fresh computation for the same target, so behavior is unchanged.
+        assertEquals(partitioning.getPartitionKey(reqA.getUri(), reqA.getVirtualHost(), null), k1,
+                "memoized key must equal a fresh computation for the current target");
+
+        // Changing the target host must invalidate the memo and yield the new host's key — otherwise a
+        // redirect could reuse a pooled connection to the wrong host.
+        Request reqB = get("http://hostb.example/").build();
+        future.setTargetRequest(reqB);
+        Object k3 = future.basePartitionKey();
+        assertNotEquals(k1, k3, "changing the target host must invalidate the memo and yield a different key");
+        assertEquals(partitioning.getPartitionKey(reqB.getUri(), reqB.getVirtualHost(), null), k3,
+                "after setTargetRequest the key must match a fresh computation for the new target");
+    }
 }