Change Body.transferTo to take a ByteBuf instead of a ByteBuffer, close AsyncHttpClient#1020

Author: slandelle

client/src/main/java/org/asynchttpclient/netty/request/body/BodyChunkedInput.java

@@ -13,15 +13,12 @@
  */
 package org.asynchttpclient.netty.request.body;
 
-import static org.asynchttpclient.util.Assertions.*;
-import org.asynchttpclient.request.body.Body;
-
+import static org.asynchttpclient.util.Assertions.assertNotNull;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.stream.ChunkedInput;
 
-import java.nio.ByteBuffer;
+import org.asynchttpclient.request.body.Body;
 
 /**
  * Adapts a {@link Body} to Netty's {@link ChunkedInput}.
@@ -52,20 +49,17 @@ public ByteBuf readChunk(ChannelHandlerContext ctx) throws Exception {
         if (endOfInput)
             return null;
 
-        // FIXME pass a visitor so we can directly pass a pooled ByteBuf
-        ByteBuffer buffer = ByteBuffer.allocate(chunkSize);
+        ByteBuf buffer = ctx.alloc().buffer(chunkSize);
         Body.BodyState state = body.transferTo(buffer);
         switch (state) {
             case STOP:
                 endOfInput = true;
-                buffer.flip();
-                return Unpooled.wrappedBuffer(buffer);
+                return buffer;
             case SUSPEND:
                 //this will suspend the stream in ChunkedWriteHandler
                 return null;
             case CONTINUE:
-                buffer.flip();
-                return Unpooled.wrappedBuffer(buffer);
+                return buffer;
             default:
                 throw new IllegalStateException("Unknown state: " + state);
         }

client/src/main/java/org/asynchttpclient/request/body/Body.java

@@ -13,9 +13,10 @@
 
 package org.asynchttpclient.request.body;
 
+import io.netty.buffer.ByteBuf;
+
 import java.io.Closeable;
 import java.io.IOException;
-import java.nio.ByteBuffer;
 
 /**
  * A request body.
@@ -50,10 +51,9 @@ enum BodyState {
     /**
      * Reads the next chunk of bytes from the body.
      *
-     * @param buffer The buffer to store the chunk in, must not be {@code null}.
+     * @param target The buffer to store the chunk in, must not be {@code null}.
      * @return The non-negative number of bytes actually read or {@code -1} if the body has been read completely.
      * @throws IOException If the chunk could not be read.
      */
-    // FIXME introduce a visitor pattern so that Netty can pass a pooled buffer
-    BodyState transferTo(ByteBuffer buffer) throws IOException;
+    BodyState transferTo(ByteBuf target) throws IOException;
 }

client/src/main/java/org/asynchttpclient/request/body/generator/ByteArrayBodyGenerator.java

@@ -12,8 +12,9 @@
  */
 package org.asynchttpclient.request.body.generator;
 
+import io.netty.buffer.ByteBuf;
+
 import java.io.IOException;
-import java.nio.ByteBuffer;
 
 import org.asynchttpclient.request.body.Body;
 
@@ -36,19 +37,20 @@ public long getContentLength() {
             return bytes.length;
         }
 
-        public BodyState transferTo(ByteBuffer byteBuffer) throws IOException {
+        public BodyState transferTo(ByteBuf target) throws IOException {
 
             if (eof) {
                 return BodyState.STOP;
             }
 
             final int remaining = bytes.length - lastPosition;
-            if (remaining <= byteBuffer.capacity()) {
-                byteBuffer.put(bytes, lastPosition, remaining);
+            final int initialTargetWritableBytes = target.writableBytes();
+            if (remaining <= initialTargetWritableBytes) {
+                target.writeBytes(bytes, lastPosition, remaining);
                 eof = true;
             } else {
-                byteBuffer.put(bytes, lastPosition, byteBuffer.capacity());
-                lastPosition = lastPosition + byteBuffer.capacity();
+                target.writeBytes(bytes, lastPosition, initialTargetWritableBytes);
+                lastPosition += initialTargetWritableBytes;
             }
             return BodyState.CONTINUE;
         }

client/src/main/java/org/asynchttpclient/request/body/generator/InputStreamBodyGenerator.java

@@ -13,13 +13,14 @@
 
 package org.asynchttpclient.request.body.generator;
 
-import org.asynchttpclient.request.body.Body;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import io.netty.buffer.ByteBuf;
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.nio.ByteBuffer;
+
+import org.asynchttpclient.request.body.Body;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A {@link BodyGenerator} which use an {@link InputStream} for reading bytes, without having to read the entire stream in memory.
@@ -61,10 +62,10 @@ public long getContentLength() {
             return -1L;
         }
 
-        public BodyState transferTo(ByteBuffer buffer) throws IOException {
+        public BodyState transferTo(ByteBuf target) throws IOException {
 
             // To be safe.
-            chunk = new byte[buffer.remaining() - 10];
+            chunk = new byte[target.writableBytes() - 10];
 
             int read = -1;
             boolean write = false;
@@ -75,7 +76,7 @@ public BodyState transferTo(ByteBuffer buffer) throws IOException {
             }
 
             if (read > 0) {
-                buffer.put(chunk, 0, read);
+                target.writeBytes(chunk, 0, read);
                 write = true;
             }
             return write ? BodyState.CONTINUE : BodyState.STOP;

client/src/main/java/org/asynchttpclient/request/body/generator/ReactiveStreamsBodyGenerator.java

@@ -12,6 +12,8 @@
  */
 package org.asynchttpclient.request.body.generator;
 
+import io.netty.buffer.ByteBuf;
+
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -78,11 +80,11 @@ public long getContentLength() {
         }
 
         @Override
-        public BodyState transferTo(ByteBuffer buffer) throws IOException {
+        public BodyState transferTo(ByteBuf target) throws IOException {
             if(initialized.compareAndSet(false, true))
                 publisher.subscribe(subscriber);
 
-            return body.transferTo(buffer);
+            return body.transferTo(target);
         }
     }
 

client/src/main/java/org/asynchttpclient/request/body/generator/SimpleFeedableBodyGenerator.java

@@ -13,6 +13,8 @@
  */
 package org.asynchttpclient.request.body.generator;
 
+import io.netty.buffer.ByteBuf;
+
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Queue;
@@ -52,20 +54,20 @@ public long getContentLength() {
         }
 
         @Override
-        public BodyState transferTo(final ByteBuffer buffer) throws IOException {
+        public BodyState transferTo(final ByteBuf target) throws IOException {
             switch (state) {
                 case CONTINUE:
-                    return readNextPart(buffer);
+                    return readNextPart(target);
                 case STOP:
                     return BodyState.STOP;
                 default:
                     throw new IllegalStateException("Illegal process state.");
             }
         }
 
-        private BodyState readNextPart(ByteBuffer buffer) throws IOException {
+        private BodyState readNextPart(ByteBuf target) throws IOException {
             BodyState res = BodyState.SUSPEND;
-            while (buffer.hasRemaining() && state != BodyState.STOP) {
+            while (target.isWritable() && state != BodyState.STOP) {
                 BodyPart nextPart = queue.peek();
                 if (nextPart == null) {
                     // Nothing in the queue. suspend stream if nothing was read. (reads == 0)
@@ -75,14 +77,14 @@ private BodyState readNextPart(ByteBuffer buffer) throws IOException {
                     queue.remove();
                 } else {
                     res = BodyState.CONTINUE;
-                    readBodyPart(buffer, nextPart);
+                    readBodyPart(target, nextPart);
                 }
             }
             return res;
         }
 
-        private void readBodyPart(ByteBuffer buffer, BodyPart part) {
-            move(buffer, part.buffer);
+        private void readBodyPart(ByteBuf target, BodyPart part) {
+            move(target, part.buffer);
 
             if (!part.buffer.hasRemaining()) {
                 if (part.isLast) {
@@ -97,12 +99,12 @@ public void close() {
         }
     }
 
-    private void move(ByteBuffer destination, ByteBuffer source) {
-        int size = Math.min(destination.remaining(), source.remaining());
+    private void move(ByteBuf target, ByteBuffer source) {
+        int size = Math.min(target.writableBytes(), source.remaining());
         if (size > 0) {
             ByteBuffer slice = source.slice();
             slice.limit(size);
-            destination.put(slice);
+            target.writeBytes(slice);
             source.position(source.position() + size);
         }
     }

client/src/main/java/org/asynchttpclient/request/body/multipart/MultipartBody.java

@@ -13,11 +13,13 @@
 package org.asynchttpclient.request.body.multipart;
 
 import static org.asynchttpclient.util.Assertions.assertNotNull;
+import static org.asynchttpclient.util.MiscUtils.closeSilently;
+import io.netty.buffer.ByteBuf;
 
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.nio.channels.WritableByteChannel;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.asynchttpclient.netty.request.body.BodyChunkedInput;
 import org.asynchttpclient.request.body.RandomAccessBody;
@@ -36,6 +38,7 @@ public class MultipartBody implements RandomAccessBody {
     private final long contentLength;
     private int currentPartIndex;
     private boolean done = false;
+    private AtomicBoolean closed = new AtomicBoolean();
 
     public MultipartBody(List<MultipartPart<? extends Part>> parts, String contentType, byte[] boundary) {
         assertNotNull(parts, "parts");
@@ -63,8 +66,10 @@ private long computeContentLength() {
     }
 
     public void close() throws IOException {
-        for (MultipartPart<? extends Part> part: parts) {
-            part.close();
+        if (closed.compareAndSet(false, true)) {
+            for (MultipartPart<? extends Part> part : parts) {
+                closeSilently(part);
+            }
         }
     }
 
@@ -81,12 +86,12 @@ public byte[] getBoundary() {
     }
 
     // Regular Body API
-    public BodyState transferTo(ByteBuffer target) throws IOException {
+    public BodyState transferTo(ByteBuf target) throws IOException {
 
         if (done)
             return BodyState.STOP;
 
-        while (target.hasRemaining() && !done) {
+        while (target.isWritable() && !done) {
             MultipartPart<? extends Part> currentPart = parts.get(currentPartIndex);
             currentPart.transferTo(target);
 

client/src/main/java/org/asynchttpclient/request/body/multipart/part/ByteArrayMultipartPart.java

@@ -1,18 +1,20 @@
 package org.asynchttpclient.request.body.multipart.part;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.nio.channels.WritableByteChannel;
 
 import org.asynchttpclient.request.body.multipart.ByteArrayPart;
 
 public class ByteArrayMultipartPart extends MultipartPart<ByteArrayPart> {
 
-    private final ByteBuffer contentBuffer;
+    private final ByteBuf contentBuffer;
 
     public ByteArrayMultipartPart(ByteArrayPart part, byte[] boundary) {
         super(part, boundary);
-        contentBuffer = ByteBuffer.wrap(part.getBytes());
+        contentBuffer = Unpooled.wrappedBuffer(part.getBytes());
     }
 
     @Override
@@ -21,7 +23,7 @@ protected long getContentLength() {
     }
 
     @Override
-    protected long transferContentTo(ByteBuffer target) throws IOException {
+    protected long transferContentTo(ByteBuf target) throws IOException {
         return transfer(contentBuffer, target, MultipartState.POST_CONTENT);
     }
 
@@ -32,5 +34,7 @@ protected long transferContentTo(WritableByteChannel target) throws IOException
 
     @Override
     public void close() {
+        super.close();
+        contentBuffer.release();
     }
 }

client/src/main/java/org/asynchttpclient/request/body/multipart/part/FileMultipartPart.java

@@ -1,9 +1,11 @@
 package org.asynchttpclient.request.body.multipart.part;
 
+import static org.asynchttpclient.util.MiscUtils.closeSilently;
+import io.netty.buffer.ByteBuf;
+
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.channels.WritableByteChannel;
 
@@ -12,7 +14,6 @@
 
 public class FileMultipartPart extends MultipartPart<FilePart> {
 
-    // FIXME make sure channel gets closed when upload crashes
     private final FileChannel channel;
     private final long length;
     private long position = 0L;
@@ -33,16 +34,16 @@ protected long getContentLength() {
     }
 
     @Override
-    protected long transferContentTo(ByteBuffer target) throws IOException {
-        int transferred = channel.read(target);
+    protected long transferContentTo(ByteBuf target) throws IOException {
+        int transferred = target.writeBytes(channel, target.writableBytes());
         position += transferred;
         if (position == length) {
             state = MultipartState.POST_CONTENT;
             channel.close();
         }
         return transferred;
     }
-    
+
     @Override
     protected long transferContentTo(WritableByteChannel target) throws IOException {
         long transferred = channel.transferTo(channel.position(), BodyChunkedInput.DEFAULT_CHUNK_SIZE, target);
@@ -55,9 +56,10 @@ protected long transferContentTo(WritableByteChannel target) throws IOException
         }
         return transferred;
     }
-    
+
     @Override
-    public void close() throws IOException {
-        channel.close();
+    public void close() {
+        super.close();
+        closeSilently(channel);
     }
 }