Bug fix for S3AsyncClient.putObject hangs if there is a connection reset while uploading of objects

Author: joviegas

.changes/next-release/bugfix-NettyNIOHttpClient-b7e834c.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Netty NIO Http Client",
+    "contributor": "",
+    "description": "Fix for Netty based client request getting stuck if connection is reset after recieveing Http Continue response."
+}

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

@@ -73,6 +73,12 @@ public final class ChannelAttributeKey {
     public static final AttributeKey<ChannelDiagnostics> CHANNEL_DIAGNOSTICS = NettyUtils.getOrCreateAttributeKey(
         "aws.http.nio.netty.async.channelDiagnostics");
 
+    /**
+     * {@link AttributeKey} to keep track of whether we have received the {@link LastHttpContent}.
+     */
+    public static final AttributeKey<Boolean> LAST_HTTP_CONTENT_RECEIVED_KEY = NettyUtils.getOrCreateAttributeKey(
+        "aws.http.nio.netty.async.lastHttpContentReceived");
+
     /**
      * {@link AttributeKey} to keep track of whether we should close the connection after this request
      * has completed.
@@ -100,12 +106,6 @@ public final class ChannelAttributeKey {
     static final AttributeKey<Long> RESPONSE_DATA_READ = NettyUtils.getOrCreateAttributeKey(
         "aws.http.nio.netty.async.responseDataRead");
 
-    /**
-     * {@link AttributeKey} to keep track of whether we have received the {@link LastHttpContent}.
-     */
-    static final AttributeKey<Boolean> LAST_HTTP_CONTENT_RECEIVED_KEY = NettyUtils.getOrCreateAttributeKey(
-        "aws.http.nio.netty.async.lastHttpContentReceived");
-
     static final AttributeKey<CompletableFuture<Void>> EXECUTE_FUTURE_KEY = NettyUtils.getOrCreateAttributeKey(
             "aws.http.nio.netty.async.executeFuture");
 

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

@@ -49,7 +49,6 @@ public void channelReleased(Channel channel) {
         if (channel.isOpen() || channel.isRegistered()) {
             removeIfExists(channel.pipeline(),
                            HttpStreamsClientHandler.class,
-                           LastHttpContentHandler.class,
                            FlushOnReadHandler.class,
                            ResponseHandler.class,
                            ReadTimeoutHandler.class,

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

@@ -220,7 +220,6 @@ private void configurePipeline() throws IOException {
                 throw new IOException("Unknown protocol: " + protocol);
         }
 
-        pipeline.addLast(LastHttpContentHandler.create());
         if (protocol == Protocol.HTTP2) {
             pipeline.addLast(FlushOnReadHandler.getInstance());
         }

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

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.http.nio.netty.internal.nrs;
 
+import static software.amazon.awssdk.http.nio.netty.internal.ChannelAttributeKey.LAST_HTTP_CONTENT_RECEIVED_KEY;
+
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelFutureListener;
@@ -211,7 +213,7 @@ private void handleCancelled(ChannelHandlerContext ctx, InT msg) {
     private void handleReadHttpContent(ChannelHandlerContext ctx, HttpContent content) {
         if (!ignoreBodyRead) {
             if (content instanceof LastHttpContent) {
-
+                ctx.channel().attr(LAST_HTTP_CONTENT_RECEIVED_KEY).set(true);
                 if (content.content().readableBytes() > 0 ||
                         !((LastHttpContent) content).trailingHeaders().isEmpty()) {
                     // It has data or trailing headers, send them

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

@@ -0,0 +1,56 @@
+package software.amazon.awssdk.http.nio.netty.fault;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
+
+public class SdkTestHttpContentPublisher implements SdkHttpContentPublisher {
+    private final byte[] body;
+    private final AtomicReference<Subscriber<? super ByteBuffer>> subscriber = new AtomicReference<>(null);
+    private final AtomicBoolean complete = new AtomicBoolean(false);
+
+    public SdkTestHttpContentPublisher(byte[] body) {
+        this.body = body;
+    }
+
+    @Override
+    public void subscribe(Subscriber<? super ByteBuffer> s) {
+        boolean wasFirstSubscriber = subscriber.compareAndSet(null, s);
+
+        SdkTestHttpContentPublisher publisher = this;
+
+        if (wasFirstSubscriber) {
+            s.onSubscribe(new Subscription() {
+                @Override
+                public void request(long n) {
+                    publisher.request(n);
+                }
+
+                @Override
+                public void cancel() {
+                    // Do nothing
+                }
+            });
+        } else {
+            s.onError(new RuntimeException("Only allow one subscriber"));
+        }
+    }
+
+    protected void request(long n) {
+        // Send the whole body if they request >0 ByteBuffers
+        if (n >  0 && !complete.get()) {
+            complete.set(true);
+            subscriber.get().onNext(ByteBuffer.wrap(body));
+            subscriber.get().onComplete();
+        }
+    }
+
+    @Override
+    public Optional<Long> contentLength() {
+        return Optional.of((long)body.length);
+    }
+}