[netty-nio-client] Ensure initial channel used for protocol detection is released before re-acquiring (#2882)

* [netty-nio-client] Ensure initial channel used for protocol detection is released before re-acquiring

HttpOrHttp2ChannelPool establishes an initial connection to determine if
H1 or H2 should be used for the ChannelPool. It then releases the Channel
to the Pool to allow it to be re-acquired and reused. However, in some
(perhaps rare) circumstances, it's possible for a subsequent acquire to race
against the release, causing the second acquire to establish a new channel
if the release had not yet completed. This may add unnecessary delay and
incur unnecessary channel overhead. Delay because the release is likely to
complete in sub-millisecond time while a new connection handshake may take
many milliseconds to complete. Overhead because a second channel may linger
in the pool even if the user only ever uses one channel at a time.

Author: Bennett-Lynch

.changes/next-release/bugfix-nettynioclient-0fb07b2.json

@@ -0,0 +1,6 @@
+{
+    "category": "Netty NIO HTTP Client", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Ensure initial channel used for protocol detection is released before re-acquiring"
+}

http-clients/netty-nio-client/src/main/java/software/amazon/awssdk/http/nio/netty/internal/http2/HttpOrHttp2ChannelPool.java

@@ -17,6 +17,7 @@
 
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.PROTOCOL_FUTURE;
 import static software.amazon.awssdk.http.nio.netty.internal.utils.NettyUtils.doInEventLoop;
+import static software.amazon.awssdk.http.nio.netty.internal.utils.NettyUtils.runOrPropagate;
 
 import io.netty.channel.Channel;
 import io.netty.channel.EventLoop;
@@ -140,7 +141,7 @@ private void completeProtocolConfiguration(Channel newChannel, Protocol protocol
                 closeAndRelease(newChannel, new IllegalStateException("Pool closed"));
             } else {
                 try {
-                    protocolImplPromise.setSuccess(configureProtocol(newChannel, protocol));
+                    configureProtocol(newChannel, protocol);
                 } catch (Throwable e) {
                     closeAndRelease(newChannel, e);
                 }
@@ -154,7 +155,7 @@ private void closeAndRelease(Channel newChannel, Throwable e) {
         protocolImplPromise.setFailure(e);
     }
 
-    private ChannelPool configureProtocol(Channel newChannel, Protocol protocol) {
+    private void configureProtocol(Channel newChannel, Protocol protocol) {
         if (Protocol.HTTP1_1 == protocol) {
             // For HTTP/1.1 we use a traditional channel pool without multiplexing
             SdkChannelPool idleConnectionMetricChannelPool = new IdleConnectionCountingChannelPool(eventLoop, delegatePool);
@@ -180,8 +181,10 @@ private ChannelPool configureProtocol(Channel newChannel, Protocol protocol) {
                                                  .build();
         }
         // Give the channel back so it can be acquired again by protocolImpl
-        delegatePool.release(newChannel);
-        return protocolImpl;
+        // Await the release completion to ensure we do not unnecessarily acquire a second channel
+        delegatePool.release(newChannel).addListener(runOrPropagate(protocolImplPromise, () -> {
+            protocolImplPromise.trySuccess(protocolImpl);
+        }));
     }
 
     @Override

http-clients/netty-nio-client/src/main/java/software/amazon/awssdk/http/nio/netty/internal/utils/NettyUtils.java

@@ -29,6 +29,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLParameters;
@@ -221,4 +222,48 @@ private static void configureSslEngine(SSLEngine sslEngine) {
         sslParameters.setEndpointIdentificationAlgorithm("HTTPS");
         sslEngine.setSSLParameters(sslParameters);
     }
+
+    /**
+     * Create a {@link GenericFutureListener} that will propagate any failures or cancellations to the provided {@link Promise},
+     * or invoke the provided {@link Consumer} with the result of a successful operation completion. This is useful for chaining
+     * together multiple futures that may depend upon each other but that may not have the same return type.
+     * <p>
+     * Note that if you do not need the value returned by a successful completion (or if it returns {@link Void}) you may use
+     * {@link #runOrPropagate(Promise, Runnable)} instead.
+     *
+     * @param destination the Promise to notify upon failure or cancellation
+     * @param onSuccess   the Consumer to invoke upon success
+     */
+    public static <T> GenericFutureListener<Future<T>> consumeOrPropagate(Promise<?> destination, Consumer<T> onSuccess) {
+        return f -> {
+            if (f.isSuccess()) {
+                T result = f.getNow();
+                onSuccess.accept(result);
+            } else if (f.isCancelled()) {
+                destination.cancel(false);
+            } else {
+                destination.tryFailure(f.cause());
+            }
+        };
+    }
+
+    /**
+     * Create a {@link GenericFutureListener} that will propagate any failures or cancellations to the provided {@link Promise},
+     * or invoke the provided {@link Runnable} upon successful operation completion. This is useful for chaining together multiple
+     * futures that may depend upon each other but that may not have the same return type.
+     *
+     * @param destination the Promise to notify upon failure or cancellation
+     * @param onSuccess   the Runnable to invoke upon success
+     */
+    public static <T> GenericFutureListener<Future<T>> runOrPropagate(Promise<?> destination, Runnable onSuccess) {
+        return f -> {
+            if (f.isSuccess()) {
+                onSuccess.run();
+            } else if (f.isCancelled()) {
+                destination.cancel(false);
+            } else {
+                destination.tryFailure(f.cause());
+            }
+        };
+    }
 }

http-clients/netty-nio-client/src/test/java/software/amazon/awssdk/http/nio/netty/internal/http2/HttpOrHttp2ChannelPoolTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.http.nio.netty.internal.http2;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -241,7 +242,9 @@ public void incompleteProtocolFutureDelaysMetricsDelegationAndForwardsSuccessFor
 
     public void incompleteProtocolFutureDelaysMetricsDelegationAndForwardsSuccessForProtocol(Protocol protocol) throws Exception {
         Promise<Channel> acquirePromise = eventLoopGroup.next().newPromise();
+        Promise<Void> releasePromise = eventLoopGroup.next().newPromise();
         when(mockDelegatePool.acquire()).thenReturn(acquirePromise);
+        when(mockDelegatePool.release(any(Channel.class))).thenReturn(releasePromise);
 
         // startConnection
         httpOrHttp2ChannelPool.acquire();
@@ -258,6 +261,7 @@ public void incompleteProtocolFutureDelaysMetricsDelegationAndForwardsSuccessFor
         eventLoopGroup.register(channel);
         channel.attr(PROTOCOL_FUTURE).set(CompletableFuture.completedFuture(protocol));
         acquirePromise.setSuccess(channel);
+        releasePromise.setSuccess(null);
 
         metricsFuture.join();
         MetricCollection metrics = metricCollector.collect();
@@ -267,4 +271,31 @@ public void incompleteProtocolFutureDelaysMetricsDelegationAndForwardsSuccessFor
         assertThat(metrics.metricValues(HttpMetric.AVAILABLE_CONCURRENCY).get(0)).isBetween(0, 1);
         assertThat(metrics.metricValues(HttpMetric.LEASED_CONCURRENCY).get(0)).isBetween(0, 1);
     }
+
+    @Test(timeout = 5_000)
+    public void protocolFutureAwaitsReleaseFuture() throws Exception {
+        Promise<Channel> delegateAcquirePromise = eventLoopGroup.next().newPromise();
+        Promise<Void> releasePromise = eventLoopGroup.next().newPromise();
+        when(mockDelegatePool.acquire()).thenReturn(delegateAcquirePromise);
+        when(mockDelegatePool.release(any(Channel.class))).thenReturn(releasePromise);
+
+        MockChannel channel = new MockChannel();
+        eventLoopGroup.register(channel);
+        channel.attr(PROTOCOL_FUTURE).set(CompletableFuture.completedFuture(Protocol.HTTP1_1));
+
+        // Acquire a new connection and save the returned future
+        Future<Channel> acquireFuture = httpOrHttp2ChannelPool.acquire();
+        
+        // Return a successful connection from the delegate pool
+        delegateAcquirePromise.setSuccess(channel);
+
+        // The returned future should not complete until the release completes
+        assertThat(acquireFuture.isDone()).isFalse();
+        
+        // Complete the release
+        releasePromise.setSuccess(null);
+        
+        // Assert the returned future completes (within the test timeout)
+        acquireFuture.await();
+    }
 }

http-clients/netty-nio-client/src/test/java/software/amazon/awssdk/http/nio/netty/internal/utils/NettyUtilsTest.java

@@ -23,17 +23,39 @@
 import static org.mockito.Mockito.when;
 
 import io.netty.channel.Channel;
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.SslHandler;
 import io.netty.util.AttributeKey;
 import io.netty.util.concurrent.EventExecutor;
+import io.netty.util.concurrent.Future;
+import io.netty.util.concurrent.GenericFutureListener;
+import io.netty.util.concurrent.Promise;
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.net.ssl.SSLEngine;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import software.amazon.awssdk.http.nio.netty.internal.MockChannel;
 
 public class NettyUtilsTest {
+
+    private static EventLoopGroup eventLoopGroup;
+
+    @BeforeClass
+    public static void setup() {
+        eventLoopGroup = new NioEventLoopGroup(1);
+    }
+
+    @AfterClass
+    public static void teardown() throws InterruptedException {
+        eventLoopGroup.shutdownGracefully().await();
+    }
+    
     @Test
     public void testGetOrCreateAttributeKey_calledTwiceWithSameName_returnsSameInstance() {
         String attr = "NettyUtilsTest.Foo";
@@ -77,4 +99,98 @@ public void doInEventLoop_notInEventLoop_submits() {
         NettyUtils.doInEventLoop(mockExecutor, () -> {});
         verify(mockExecutor).submit(any(Runnable.class));
     }
+
+    @Test
+    public void runOrPropagate_success_runs() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+        AtomicBoolean reference = new AtomicBoolean();
+
+        GenericFutureListener<Future<Void>> listener =
+            NettyUtils.runOrPropagate(destination, () -> reference.set(true));
+
+        Promise<Void> source = eventLoopGroup.next().newPromise();
+        source.setSuccess(null);
+        listener.operationComplete(source);
+
+        assertThat(reference.get()).isTrue();
+    }
+
+    @Test
+    public void runOrPropagate_exception_propagates() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+
+        GenericFutureListener<Future<Void>> listener =
+            NettyUtils.runOrPropagate(destination, () -> {
+            });
+
+        Promise<Void> source = eventLoopGroup.next().newPromise();
+        source.setFailure(new RuntimeException("Intentional exception for testing purposes"));
+        listener.operationComplete(source);
+
+        assertThat(destination.cause())
+            .isInstanceOf(RuntimeException.class)
+            .hasMessage("Intentional exception for testing purposes");
+    }
+
+    @Test
+    public void runOrPropagate_cancel_propagates() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+
+        GenericFutureListener<Future<Void>> listener =
+            NettyUtils.runOrPropagate(destination, () -> {
+            });
+
+        Promise<Void> source = eventLoopGroup.next().newPromise();
+        source.cancel(false);
+        listener.operationComplete(source);
+
+        assertThat(destination.isCancelled()).isTrue();
+    }
+
+    @Test
+    public void consumeOrPropagate_success_consumes() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+        AtomicReference<String> reference = new AtomicReference<>();
+
+        GenericFutureListener<Future<String>> listener =
+            NettyUtils.consumeOrPropagate(destination, reference::set);
+
+        Promise<String> source = eventLoopGroup.next().newPromise();
+        source.setSuccess("test");
+        listener.operationComplete(source);
+
+        assertThat(reference.get()).isEqualTo("test");
+    }
+
+    @Test
+    public void consumeOrPropagate_exception_propagates() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+
+        GenericFutureListener<Future<String>> listener =
+            NettyUtils.consumeOrPropagate(destination, s -> {
+            });
+
+        Promise<String> source = eventLoopGroup.next().newPromise();
+        source.setFailure(new RuntimeException("Intentional exception for testing purposes"));
+        listener.operationComplete(source);
+
+        assertThat(destination.cause())
+            .isInstanceOf(RuntimeException.class)
+            .hasMessage("Intentional exception for testing purposes");
+    }
+
+    @Test
+    public void consumeOrPropagate_cancel_propagates() throws Exception {
+        Promise<String> destination = eventLoopGroup.next().newPromise();
+
+        GenericFutureListener<Future<String>> listener =
+            NettyUtils.consumeOrPropagate(destination, s -> {
+            });
+
+        Promise<String> source = eventLoopGroup.next().newPromise();
+        source.cancel(false);
+        listener.operationComplete(source);
+
+        assertThat(destination.isCancelled()).isTrue();
+    }
 }