[netty-nio-client] Ensure in-use channels are not incorrectly closed (#2883)

* [netty-nio-client] Ensure in-use channels are not incorrectly closed

IdleConnectionReaperHandler and OldConnectionReaperHandler are responsible for closing idle/old channels, but only when they are not marked as IN_USE. NettyRequestExecutor is responsible for marking a channel as IN_USE as part of its logic to prepare a channel for a request. However, in rare circumstances, it may be possible for a channel to reach its idle time in between being acquired from the pool and having the request initiated (especially when the channel is acquired from a thread outside the EventLoop). We can make this more resilient by more aggressively marking the channel as IN_USE as part of the channel acquire logic, and ensuring this is run before the channel's health is queried by HealthCheckedChannelPool (which will acquire a new channel, if needed).

Previously, NettyRequestExecutor was responsible for setting IN_USE and HandlerRemovingChannelPool was responsible for clearing it. This PR introduces a simple ChannelPoolListener interface that allows implementing classes to register behavior based on channels being acquired/released (while also guaranteeing the listener callbacks are executed within the channel's EventLoop). We can leverage this interface to make the IN_USE behavior more symmetrical, by using a single ChannelPoolListener to both set/clear the flag. The ChannelPoolListener interface can additionally be used to simplify some of our other ChannelPool-wrapping classes (the ones that don't need to actually manipulate the pool control flow). So this commit also takes the opportunity to migrate the existing HandlerRemovingChannelPool to leverage this interface.

Author: Bennett-Lynch

.changes/next-release/bugfix-nettynioclient-841857c.json

@@ -0,0 +1,6 @@
+{
+    "category": "Netty NIO HTTP Client", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Ensure in-use channels are not incorrectly closed"
+}

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

@@ -26,6 +26,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -236,8 +237,13 @@ private SdkChannelPool wrapBaseChannelPool(Bootstrap bootstrap, ChannelPool chan
                                                                    configuration.maxConnections(),
                                                                    configuration);
 
-        // Wrap the channel pool such that we remove request-specific handlers with each request.
-        sdkChannelPool = new HandlerRemovingChannelPool(sdkChannelPool);
+        sdkChannelPool = new ListenerInvokingChannelPool(bootstrap.config().group(), sdkChannelPool, Arrays.asList(
+            // Add a listener that ensures acquired channels are marked IN_USE and thus not eligible for certain idle timeouts.
+            InUseTrackingChannelPoolListener.create(),
+
+            // Add a listener that removes request-specific handlers with each request.
+            HandlerRemovingChannelPoolListener.create()
+        ));
 
         // Wrap the channel pool such that an individual channel can only be released to the underlying pool once.
         sdkChannelPool = new ReleaseOnceChannelPool(sdkChannelPool);

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

@@ -0,0 +1,59 @@
+/*
+ * Copyright 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;
+
+import static software.amazon.awssdk.http.nio.netty.internal.utils.ChannelUtils.removeIfExists;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandler;
+import io.netty.handler.timeout.ReadTimeoutHandler;
+import io.netty.handler.timeout.WriteTimeoutHandler;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.http.nio.netty.internal.ListenerInvokingChannelPool.ChannelPoolListener;
+import software.amazon.awssdk.http.nio.netty.internal.http2.FlushOnReadHandler;
+import software.amazon.awssdk.http.nio.netty.internal.nrs.HttpStreamsClientHandler;
+
+/**
+ * Removes any per-request {@link ChannelHandler} from the pipeline when releasing it to the pool.
+ */
+@SdkInternalApi
+public final class HandlerRemovingChannelPoolListener implements ChannelPoolListener {
+
+    private static final HandlerRemovingChannelPoolListener INSTANCE = new HandlerRemovingChannelPoolListener();
+
+    private HandlerRemovingChannelPoolListener() {
+    }
+
+    public static HandlerRemovingChannelPoolListener create() {
+        return INSTANCE;
+    }
+
+    @Override
+    public void channelReleased(Channel channel) {
+        // Only remove per request handler if the channel is registered
+        // or open since DefaultChannelPipeline would remove handlers if
+        // channel is closed and unregistered
+        if (channel.isOpen() || channel.isRegistered()) {
+            removeIfExists(channel.pipeline(),
+                           HttpStreamsClientHandler.class,
+                           LastHttpContentHandler.class,
+                           FlushOnReadHandler.class,
+                           ResponseHandler.class,
+                           ReadTimeoutHandler.class,
+                           WriteTimeoutHandler.class);
+        }
+    }
+}

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 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;
+
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.IN_USE;
+
+import io.netty.channel.Channel;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.http.nio.netty.internal.ListenerInvokingChannelPool.ChannelPoolListener;
+
+/**
+ * Marks {@link Channel}s as in-use when they are leased from the pool. An in-use channel is not eligible to be closed by {@link
+ * IdleConnectionReaperHandler} or {@link OldConnectionReaperHandler}.
+ */
+@SdkInternalApi
+public final class InUseTrackingChannelPoolListener implements ChannelPoolListener {
+
+    private static final InUseTrackingChannelPoolListener INSTANCE = new InUseTrackingChannelPoolListener();
+
+    private InUseTrackingChannelPoolListener() {
+    }
+
+    public static InUseTrackingChannelPoolListener create() {
+        return INSTANCE;
+    }
+
+    @Override
+    public void channelAcquired(Channel channel) {
+        channel.attr(IN_USE).set(true);
+    }
+
+    @Override
+    public void channelReleased(Channel channel) {
+        channel.attr(IN_USE).set(false);
+    }
+}

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

@@ -0,0 +1,139 @@
+/*
+ * Copyright 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;
+
+import static software.amazon.awssdk.http.nio.netty.internal.utils.NettyUtils.consumeOrPropagate;
+
+import io.netty.channel.Channel;
+import io.netty.channel.EventLoop;
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.pool.ChannelPoolHandler;
+import io.netty.util.concurrent.Future;
+import io.netty.util.concurrent.Promise;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.http.nio.netty.internal.http2.HttpOrHttp2ChannelPool;
+import software.amazon.awssdk.http.nio.netty.internal.utils.NettyUtils;
+import software.amazon.awssdk.metrics.MetricCollector;
+
+/**
+ * A {@link SdkChannelPool} that wraps and delegates to another {@link SdkChannelPool} while invoking {@link ChannelPoolListener}s
+ * for important events that occur.
+ * <p>
+ * {@link ChannelPoolListener} is similar to Netty's {@link ChannelPoolHandler} interface, but by invoking them as part of a
+ * {@link SdkChannelPool} wrapper, we are given more control over when they are invoked. For example, {@link
+ * HttpOrHttp2ChannelPool} may choose not to release HTTP/2 stream channels to the lowest-level pool (and instead store the
+ * channels in its own pool), but by instrumenting listeners that sit on top of this layer, we are still given visibility into
+ * these events occurring.
+ */
+@SdkInternalApi
+public final class ListenerInvokingChannelPool implements SdkChannelPool {
+    private final SdkChannelPool delegatePool;
+    private final Supplier<Promise<Channel>> promiseFactory;
+    private final List<ChannelPoolListener> listeners;
+
+    /**
+     * Listener which is called for various actions performed on a {@link SdkChannelPool}. All listener events are guaranteed to
+     * be invoked as part of the {@link Channel}'s {@link EventLoop}.
+     */
+    @SdkInternalApi
+    public interface ChannelPoolListener {
+
+        /**
+         * Called <b>after</b> a {@link Channel} was acquired by calling {@link SdkChannelPool#acquire()} or {@link
+         * SdkChannelPool#acquire(Promise)}.
+         * <p>
+         * This method will be called by the {@link EventLoop} of the {@link Channel}.
+         */
+        default void channelAcquired(Channel channel) {
+        }
+
+        /**
+         * Called <b>before</b> a {@link Channel} is released by calling {@link SdkChannelPool#release(Channel)} or {@link
+         * SdkChannelPool#release(Channel, Promise)}.
+         * <p>
+         * This method will be called by the {@link EventLoop} of the {@link Channel}.
+         */
+        default void channelReleased(Channel channel) {
+        }
+    }
+
+    public ListenerInvokingChannelPool(EventLoopGroup eventLoopGroup,
+                                       SdkChannelPool delegatePool,
+                                       List<ChannelPoolListener> listeners) {
+        this(() -> eventLoopGroup.next().newPromise(), delegatePool, listeners);
+    }
+
+    public ListenerInvokingChannelPool(Supplier<Promise<Channel>> promiseFactory,
+                                       SdkChannelPool delegatePool,
+                                       List<ChannelPoolListener> listeners) {
+        this.delegatePool = delegatePool;
+        this.promiseFactory = promiseFactory;
+        this.listeners = listeners;
+    }
+
+    @Override
+    public Future<Channel> acquire() {
+        return acquire(promiseFactory.get());
+    }
+
+    @Override
+    public Future<Channel> acquire(Promise<Channel> returnFuture) {
+        delegatePool.acquire(promiseFactory.get())
+                    .addListener(consumeOrPropagate(returnFuture, channel -> {
+                        NettyUtils.doInEventLoop(channel.eventLoop(), () -> {
+                            invokeChannelAcquired(channel);
+                            returnFuture.trySuccess(channel);
+                        }, returnFuture);
+                    }));
+        return returnFuture;
+    }
+
+    private void invokeChannelAcquired(Channel channel) {
+        listeners.forEach(listener -> listener.channelAcquired(channel));
+    }
+
+    @Override
+    public Future<Void> release(Channel channel) {
+        return release(channel, channel.eventLoop().newPromise());
+    }
+
+    @Override
+    public Future<Void> release(Channel channel, Promise<Void> returnFuture) {
+        NettyUtils.doInEventLoop(channel.eventLoop(), () -> {
+            invokeChannelReleased(channel);
+            delegatePool.release(channel, returnFuture);
+        }, returnFuture);
+        return returnFuture;
+    }
+
+    private void invokeChannelReleased(Channel channel) {
+        listeners.forEach(listener -> listener.channelReleased(channel));
+    }
+
+
+    @Override
+    public void close() {
+        delegatePool.close();
+    }
+
+    @Override
+    public CompletableFuture<Void> collectChannelPoolMetrics(MetricCollector metrics) {
+        return delegatePool.collectChannelPoolMetrics(metrics);
+    }
+}

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

@@ -192,7 +192,6 @@ private void configureChannel() {
         channel.attr(REQUEST_CONTEXT_KEY).set(context);
         channel.attr(RESPONSE_COMPLETE_KEY).set(false);
         channel.attr(LAST_HTTP_CONTENT_RECEIVED_KEY).set(false);
-        channel.attr(IN_USE).set(true);
         channel.config().setOption(ChannelOption.AUTO_READ, false);
     }