Await h2 connection releases before closing pool

.changes/next-release/bugfix-NettyNIOHTTPClient-0442e07.json

@@ -0,0 +1,5 @@
+{
+    "category": "Netty NIO HTTP Client", 
+    "type": "bugfix", 
+    "description": "Fix a bug where it's possible for an HTTP2 channel pool to be closed while some channels are still being released causing them to be left open and leaked."
+}

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

@@ -26,7 +26,10 @@
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.Promise;
 import java.util.ArrayList;
+import java.util.Collection;
+
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.http.nio.netty.internal.utils.BetterFixedChannelPool;
 
 /**
@@ -48,6 +51,7 @@ public class Http2MultiplexedChannelPool implements ChannelPool {
     private final ChannelPool connectionPool;
     private final long maxConcurrencyPerConnection;
     private final ArrayList<MultiplexedChannelRecord> connections;
+    private boolean closed = false;
 
     /**
      * @param connectionPool Connection pool for parent channels (i.e. the socket channel).
@@ -65,6 +69,17 @@ public class Http2MultiplexedChannelPool implements ChannelPool {
         this.connections = new ArrayList<>();
     }
 
+    @SdkTestInternalApi
+    Http2MultiplexedChannelPool(ChannelPool connectionPool,
+                                EventLoop eventLoop,
+                                long maxConcurrencyPerConnection,
+                                Collection<MultiplexedChannelRecord> connections) {
+        this.connectionPool = connectionPool;
+        this.eventLoop = eventLoop;
+        this.maxConcurrencyPerConnection = maxConcurrencyPerConnection;
+        this.connections = new ArrayList<>(connections);
+    }
+
     @Override
     public Future<Channel> acquire() {
         return acquire(new DefaultPromise<>(eventLoop));
@@ -77,6 +92,10 @@ public Future<Channel> acquire(Promise<Channel> promise) {
     }
 
     private Future<Channel> acquire0(Promise<Channel> promise) {
+        if (closed) {
+            return promise.setFailure(new IllegalStateException("Channel pool is closed!"));
+        }
+
         for (MultiplexedChannelRecord connection : connections) {
             if (connection.availableStreams() > 0) {
                 connection.acquire(promise);
@@ -150,7 +169,27 @@ private void releaseParentChannel(Channel parentChannel) {
 
     @Override
     public void close() {
-        doInEventLoop(eventLoop, connectionPool::close);
+        try {
+            setClosedFlag().await();
+            for (MultiplexedChannelRecord c : connections) {
+                Future<Channel> f = c.getConnectionFuture();
+                f.await();
+                if (f.isSuccess()) {
+                    connectionPool.release(f.getNow()).await();
+                }
+            }
+            connectionPool.close();
+        } catch (InterruptedException ie) {
+            throw new RuntimeException(ie);
+        }
     }
 
+    private Promise<Void> setClosedFlag() {
+        Promise<Void> closedFuture = eventLoop.newPromise();
+        doInEventLoop(eventLoop, () -> {
+            closed = true;
+            closedFuture.setSuccess(null);
+        });
+        return closedFuture;
+    }
 }

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

@@ -33,7 +33,9 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
+
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.http.Protocol;
 
 /**
@@ -64,6 +66,18 @@ public final class MultiplexedChannelRecord {
         this.channelReleaser = channelReleaser;
     }
 
+    @SdkTestInternalApi
+    MultiplexedChannelRecord(Future<Channel> connectionFuture,
+                             Channel connection,
+                             long maxConcurrencyPerConnection,
+                             BiConsumer<Channel, MultiplexedChannelRecord> channelReleaser) {
+        this.connectionFuture = connectionFuture;
+        this.childChannels = new ConcurrentHashMap<>(saturatedCast(maxConcurrencyPerConnection));
+        this.availableStreams = new AtomicLong(maxConcurrencyPerConnection);
+        this.channelReleaser = channelReleaser;
+        this.connection = connection;
+    }
+
     MultiplexedChannelRecord acquire(Promise<Channel> channelPromise) {
         availableStreams.decrementAndGet();
         if (connection != null) {
@@ -140,6 +154,10 @@ void release(Channel channel) {
         childChannels.remove(channel.id());
     }
 
+    public Future<Channel> getConnectionFuture() {
+        return connectionFuture;
+    }
+
     long availableStreams() {
         return availableStreams.get();
     }

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

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.http.nio.netty.internal.http2;
+
+import io.netty.channel.Channel;
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.nio.NioEventLoopGroup;
+import io.netty.channel.pool.ChannelPool;
+import io.netty.channel.socket.SocketChannel;
+import io.netty.channel.socket.nio.NioSocketChannel;
+import io.netty.util.concurrent.DefaultPromise;
+import io.netty.util.concurrent.Promise;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+import org.mockito.runners.MockitoJUnitRunner;
+
+import java.util.Collections;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+
+/**
+ * Tests for {@link Http2MultiplexedChannelPool}.
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class Http2MultiplexedChannelPoolTest {
+    private static EventLoopGroup loopGroup;
+
+    @BeforeClass
+    public static void setup() {
+        loopGroup = new NioEventLoopGroup(4);
+    }
+
+    @AfterClass
+    public static void teardown() {
+        loopGroup.shutdownGracefully().awaitUninterruptibly();
+    }
+
+    @Test
+    public void closeWaitsForConnectionToBeReleasedBeforeClosingConnectionPool() throws InterruptedException {
+        SocketChannel channel = new NioSocketChannel();
+        try {
+            loopGroup.register(channel).awaitUninterruptibly();
+            Promise<Channel> channelPromise = new DefaultPromise<>(loopGroup.next());
+            channelPromise.setSuccess(channel);
+
+            ChannelPool connectionPool = Mockito.mock(ChannelPool.class);
+            Promise<Void> releasePromise = Mockito.spy(new DefaultPromise<>(loopGroup.next()));
+            Mockito.doCallRealMethod().when(releasePromise).await();
+            releasePromise.setSuccess(null);
+            Mockito.when(connectionPool.release(Mockito.eq(channel))).thenReturn(releasePromise);
+
+            MultiplexedChannelRecord record = new MultiplexedChannelRecord(channelPromise,
+                                                                           channel,
+                                                                           8,
+                                                                           (ch, rec) -> {});
+            Http2MultiplexedChannelPool h2Pool = new Http2MultiplexedChannelPool(connectionPool, loopGroup.next(), 2, Collections.singletonList(record));
+
+            h2Pool.close();
+
+            InOrder inOrder = Mockito.inOrder(connectionPool, releasePromise);
+            inOrder.verify(releasePromise).await();
+            inOrder.verify(connectionPool).close();
+        } finally {
+            channel.close().awaitUninterruptibly();
+        }
+    }
+
+    @Test
+    public void acquireAfterCloseFails() throws InterruptedException {
+        ChannelPool connectionPool = Mockito.mock(ChannelPool.class);
+
+        Http2MultiplexedChannelPool h2Pool = new Http2MultiplexedChannelPool(connectionPool, loopGroup.next(), 2, Collections.emptyList());
+
+        h2Pool.close();
+
+        assertThat(h2Pool.acquire().await().isSuccess()).isFalse();
+    }
+
+    @Test
+    public void releaseAfterCloseFails() throws InterruptedException {
+        ChannelPool connectionPool = Mockito.mock(ChannelPool.class);
+
+        Http2MultiplexedChannelPool h2Pool = new Http2MultiplexedChannelPool(connectionPool, loopGroup.next(), 2, Collections.emptyList());
+
+        h2Pool.close();
+
+        assertThat(h2Pool.release(Mockito.mock(Channel.class)).await().isSuccess()).isFalse();
+    }
+}