Extend Response to support returning ByteBufs (#1954)

NettyResponse is already dealing with ByteBufs under the hood, this
lets the caller access them cleanly. This allows callers that are
already using Netty to use AHC much more efficiently.

Fixes #1953

Co-authored-by: David Golombek <[email protected]>

Author: Dave Golombek

client/src/main/java/org/asynchttpclient/HttpResponseBodyPart.java

@@ -15,6 +15,7 @@
  */
 package org.asynchttpclient;
 
+import io.netty.buffer.ByteBuf;
 import java.nio.ByteBuffer;
 
 /**
@@ -44,6 +45,12 @@ protected HttpResponseBodyPart(boolean last) {
      */
     public abstract ByteBuffer getBodyByteBuffer();
 
+    /**
+     * @return the {@link ByteBuf} of the bytes read from the response's chunk.
+     * The {@link ByteBuf}'s capacity is equal to the number of bytes available.
+     */
+    public abstract ByteBuf getBodyByteBuf();
+
     /**
      * @return true if this is the last part.
      */

client/src/main/java/org/asynchttpclient/Response.java

@@ -16,6 +16,7 @@
  */
 package org.asynchttpclient;
 
+import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.cookie.Cookie;
 import org.asynchttpclient.netty.NettyResponse;
@@ -61,6 +62,13 @@ public interface Response {
      */
     ByteBuffer getResponseBodyAsByteBuffer();
 
+    /**
+     * Return the entire response body as a ByteBuf.
+     *
+     * @return the entire response body as a ByteBuf.
+     */
+    ByteBuf getResponseBodyAsByteBuf();
+
     /**
      * Returns an input stream for the response body. Note that you should not try to get this more than once, and that you should not close the stream.
      *

client/src/main/java/org/asynchttpclient/netty/EagerResponseBodyPart.java

@@ -17,6 +17,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufUtil;
+import io.netty.buffer.Unpooled;
 import org.asynchttpclient.HttpResponseBodyPart;
 
 import java.nio.ByteBuffer;
@@ -53,4 +54,9 @@ public int length() {
     public ByteBuffer getBodyByteBuffer() {
         return ByteBuffer.wrap(bytes);
     }
+
+    @Override
+    public ByteBuf getBodyByteBuf() {
+        return Unpooled.wrappedBuffer(bytes);
+    }
 }

client/src/main/java/org/asynchttpclient/netty/LazyResponseBodyPart.java

@@ -33,7 +33,8 @@ public LazyResponseBodyPart(ByteBuf buf, boolean last) {
         this.buf = buf;
     }
 
-    public ByteBuf getBuf() {
+    @Override
+    public ByteBuf getBodyByteBuf() {
         return buf;
     }
 

client/src/main/java/org/asynchttpclient/netty/NettyResponse.java

@@ -15,6 +15,9 @@
  */
 package org.asynchttpclient.netty;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.CompositeByteBuf;
 import io.netty.handler.codec.http.EmptyHttpHeaders;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.cookie.ClientCookieDecoder;
@@ -192,6 +195,15 @@ public ByteBuffer getResponseBodyAsByteBuffer() {
         return target;
     }
 
+    @Override
+    public ByteBuf getResponseBodyAsByteBuf() {
+        CompositeByteBuf compositeByteBuf = ByteBufAllocator.DEFAULT.compositeBuffer(bodyParts.size());
+        for (HttpResponseBodyPart part : bodyParts) {
+            compositeByteBuf.addComponent(true, part.getBodyByteBuf());
+        }
+        return compositeByteBuf;
+    }
+
     @Override
     public String getResponseBody() {
         return getResponseBody(withDefault(extractContentTypeCharsetAttribute(getContentType()), UTF_8));

client/src/test/java/org/asynchttpclient/netty/NettyAsyncResponseTest.java

@@ -13,15 +13,16 @@
 package org.asynchttpclient.netty;
 
 import io.github.artsok.RepeatedIfExceptionsTest;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.DefaultHttpHeaders;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.cookie.Cookie;
+import org.asynchttpclient.HttpResponseBodyPart;
 
+import java.nio.charset.StandardCharsets;
 import java.text.SimpleDateFormat;
-import java.util.Date;
-import java.util.List;
-import java.util.Locale;
-import java.util.TimeZone;
+import java.util.*;
 
 import static io.netty.handler.codec.http.HttpHeaderNames.SET_COOKIE;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -73,4 +74,16 @@ public void testCookieParseWeirdExpiresValue() {
         Cookie cookie = cookies.get(0);
         assertEquals(Long.MIN_VALUE, cookie.maxAge());
     }
+
+    @RepeatedIfExceptionsTest(repeats = 5)
+    public void testGetResponseBodyAsByteBuffer() {
+        List<HttpResponseBodyPart> bodyParts = new LinkedList<>();
+        bodyParts.add(new LazyResponseBodyPart(Unpooled.wrappedBuffer("Hello ".getBytes()), false));
+        bodyParts.add(new LazyResponseBodyPart(Unpooled.wrappedBuffer("World".getBytes()), true));
+        NettyResponse response = new NettyResponse(new NettyResponseStatus(null, null, null), null, bodyParts);
+
+        ByteBuf body = response.getResponseBodyAsByteBuf();
+        assertEquals("Hello World", body.toString(StandardCharsets.UTF_8));
+        body.release();
+    }
 }