Add LastHttpContentHandler to flag if the LastHttpContent has been received

Set KEEP_ALIVE to false if the channel is to be closed

Author: zoewangg

.changes/next-release/bugfix-NettyNIOHttpClient-5af2afd.json

@@ -0,0 +1,5 @@
+{
+    "category": "Netty NIO Http Client", 
+    "type": "bugfix", 
+    "description": "Fix a race condition where the channel is closed right after all content is buffered, causing `server failed to complete the response` error by adding a flag when `LastHttpContentHandler` is received."
+}

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

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.http.nio.netty.internal;
 
 import io.netty.channel.Channel;
+import io.netty.handler.codec.http.LastHttpContent;
 import io.netty.util.AttributeKey;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
@@ -60,6 +61,12 @@ public final class ChannelAttributeKey {
     static final AttributeKey<Boolean> RESPONSE_COMPLETE_KEY = AttributeKey.newInstance(
         "aws.http.nio.netty.async.responseComplete");
 
+    /**
+     * {@link AttributeKey} to keep track of whether we have received the {@link LastHttpContent}.
+     */
+    static final AttributeKey<Boolean> LAST_HTTP_CONTENT_RECEIVED_KEY = AttributeKey.newInstance(
+        "aws.http.nio.netty.async.lastHttpContentReceived");
+
     static final AttributeKey<CompletableFuture<Void>> EXECUTE_FUTURE_KEY = AttributeKey.newInstance(
             "aws.http.nio.netty.async.executeFuture");
 
@@ -72,7 +79,6 @@ public final class ChannelAttributeKey {
      */
     static final AttributeKey<Boolean> KEEP_ALIVE = AttributeKey.newInstance("aws.http.nio.netty.async.keepAlive");
 
-
     /**
      * Whether the channel is still in use
      */

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

@@ -78,6 +78,7 @@ private void removePerRequestHandlers(Channel channel) {
         if (channel.isOpen() || channel.isRegistered()) {
             removeIfExists(channel.pipeline(),
                            HttpStreamsClientHandler.class,
+                           LastHttpContentHandler.class,
                            ResponseHandler.class,
                            ReadTimeoutHandler.class,
                            WriteTimeoutHandler.class);

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2019 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.LAST_HTTP_CONTENT_RECEIVED_KEY;
+
+import io.netty.channel.ChannelHandler;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.handler.codec.http.LastHttpContent;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * Marks {@code ChannelAttributeKey.LAST_HTTP_CONTENT_RECEIVED_KEY} if {@link LastHttpContent} is received.
+ */
+@SdkInternalApi
+@ChannelHandler.Sharable
+public final class LastHttpContentHandler extends ChannelInboundHandlerAdapter {
+    private static final LastHttpContentHandler INSTANCE = new LastHttpContentHandler();
+    private static final Logger logger = Logger.loggerFor(LastHttpContent.class);
+
+    @Override
+    public void channelRead(ChannelHandlerContext ctx, Object msg) {
+        if (msg instanceof LastHttpContent) {
+            logger.debug(() -> "Received LastHttpContent " + ctx.channel());
+            ctx.channel().attr(LAST_HTTP_CONTENT_RECEIVED_KEY).set(true);
+        }
+
+        ctx.fireChannelRead(msg);
+    }
+
+    public static LastHttpContentHandler create() {
+        return INSTANCE;
+    }
+}

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

@@ -20,6 +20,8 @@
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.EXECUTE_FUTURE_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.EXECUTION_ID_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.IN_USE;
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.KEEP_ALIVE;
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.LAST_HTTP_CONTENT_RECEIVED_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.REQUEST_CONTEXT_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.RESPONSE_COMPLETE_KEY;
 
@@ -145,6 +147,7 @@ private void configureChannel() {
         channel.attr(EXECUTE_FUTURE_KEY).set(executeFuture);
         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);
     }
@@ -162,6 +165,7 @@ private boolean tryConfigurePipeline() {
             return false;
         }
 
+        pipeline.addLast(LastHttpContentHandler.create());
         pipeline.addLast(new HttpStreamsClientHandler());
         pipeline.addLast(ResponseHandler.getInstance());
 
@@ -336,6 +340,7 @@ private String getMessageForClosedChannel() {
      */
     private void closeAndRelease(Channel channel) {
         log.trace("closing and releasing channel {}", channel.id().asLongText());
+        channel.attr(KEEP_ALIVE).set(false);
         channel.close();
         context.channelPool().release(channel);
     }

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

@@ -19,6 +19,7 @@
 import static java.util.stream.Collectors.mapping;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.EXECUTE_FUTURE_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.KEEP_ALIVE;
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.LAST_HTTP_CONTENT_RECEIVED_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.REQUEST_CONTEXT_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.RESPONSE_COMPLETE_KEY;
 import static software.amazon.awssdk.http.nio.netty.internal.utils.ExceptionHandlingUtils.tryCatch;
@@ -138,12 +139,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
     }
 
     @Override
-    public void channelInactive(ChannelHandlerContext handlerCtx) throws Exception {
-        notifyIfResponseNotCompleted(handlerCtx);
-    }
-
-    @Override
-    public void channelUnregistered(ChannelHandlerContext handlerCtx) throws Exception {
+    public void channelInactive(ChannelHandlerContext handlerCtx) {
         notifyIfResponseNotCompleted(handlerCtx);
     }
 
@@ -158,6 +154,7 @@ public static ResponseHandler getInstance() {
      */
     private static void closeAndRelease(ChannelHandlerContext ctx) {
         Channel channel = ctx.channel();
+        channel.attr(KEEP_ALIVE).set(false);
         RequestContext requestContext = channel.attr(REQUEST_CONTEXT_KEY).get();
         ctx.close();
         requestContext.channelPool().release(channel);
@@ -386,7 +383,10 @@ private Throwable wrapException(Throwable originalCause) {
     private void notifyIfResponseNotCompleted(ChannelHandlerContext handlerCtx) {
         RequestContext requestCtx = handlerCtx.channel().attr(REQUEST_CONTEXT_KEY).get();
         boolean responseCompleted = handlerCtx.channel().attr(RESPONSE_COMPLETE_KEY).get();
-        if (!responseCompleted) {
+        boolean lastHttpContentReceived = handlerCtx.channel().attr(LAST_HTTP_CONTENT_RECEIVED_KEY).get();
+        handlerCtx.channel().attr(KEEP_ALIVE).set(false);
+
+        if (!responseCompleted && !lastHttpContentReceived) {
             IOException err = new IOException("Server failed to send complete response");
             runAndLogError("Fail to execute SdkAsyncHttpResponseHandler#onError", () -> requestCtx.handler().onError(err));
             executeFuture(handlerCtx).completeExceptionally(err);

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2010-2019 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 org.assertj.core.api.Assertions.assertThat;
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.LAST_HTTP_CONTENT_RECEIVED_KEY;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.http.LastHttpContent;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+
+public class LastHttpContentHandlerTest {
+
+    private MockChannel channel;
+    private ChannelHandlerContext handlerContext;
+    private LastHttpContentHandler contentHandler = LastHttpContentHandler.create();
+
+    @Before
+    public void setup() throws Exception {
+        channel = new MockChannel();
+        channel.attr(LAST_HTTP_CONTENT_RECEIVED_KEY).set(false);
+        handlerContext = Mockito.mock(ChannelHandlerContext.class);
+        Mockito.when(handlerContext.channel()).thenReturn(channel);
+    }
+
+    @After
+    public void cleanup() {
+        channel.close();
+    }
+
+    @Test
+    public void lastHttpContentReceived_shouldSetAttribute() {
+        LastHttpContent lastHttpContent = LastHttpContent.EMPTY_LAST_CONTENT;
+        contentHandler.channelRead(handlerContext, lastHttpContent);
+
+        assertThat(channel.attr(LAST_HTTP_CONTENT_RECEIVED_KEY).get()).isTrue();
+    }
+
+    @Test
+    public void otherContentReceived_shouldNotSetAttribute() {
+        String content = "some content";
+        contentHandler.channelRead(handlerContext, content);
+
+        assertThat(channel.attr(LAST_HTTP_CONTENT_RECEIVED_KEY).get()).isFalse();
+    }
+}

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

@@ -1,11 +1,18 @@
 package software.amazon.awssdk.http.nio.netty.internal;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import io.netty.channel.Channel;
 import io.netty.channel.EventLoop;
 import io.netty.channel.EventLoopGroup;
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.pool.ChannelPool;
 import io.netty.util.concurrent.Promise;
+import java.util.concurrent.CompletableFuture;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -14,14 +21,6 @@
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.utils.AttributeMap;
 
-import java.util.concurrent.CompletableFuture;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 public class NettyRequestExecutorTest {
 
     private ChannelPool mockChannelPool;