From e8193ae3e3c8e431d636cd4ec421ce39dc28312e Mon Sep 17 00:00:00 2001 From: Dave Golombek Date: Thu, 11 Apr 2024 03:07:25 -0400 Subject: [PATCH] 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 --- .../asynchttpclient/HttpResponseBodyPart.java | 7 +++++++ .../java/org/asynchttpclient/Response.java | 8 +++++++ .../netty/EagerResponseBodyPart.java | 6 ++++++ .../netty/LazyResponseBodyPart.java | 3 ++- .../asynchttpclient/netty/NettyResponse.java | 12 +++++++++++ .../netty/NettyAsyncResponseTest.java | 21 +++++++++++++++---- 6 files changed, 52 insertions(+), 5 deletions(-) diff --git a/client/src/main/java/org/asynchttpclient/HttpResponseBodyPart.java b/client/src/main/java/org/asynchttpclient/HttpResponseBodyPart.java index 0be4dedb5..8ab4c9bcc 100644 --- a/client/src/main/java/org/asynchttpclient/HttpResponseBodyPart.java +++ b/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. */ diff --git a/client/src/main/java/org/asynchttpclient/Response.java b/client/src/main/java/org/asynchttpclient/Response.java index 8b9c9a6f1..220d989b0 100644 --- a/client/src/main/java/org/asynchttpclient/Response.java +++ b/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. * diff --git a/client/src/main/java/org/asynchttpclient/netty/EagerResponseBodyPart.java b/client/src/main/java/org/asynchttpclient/netty/EagerResponseBodyPart.java index 44045e80b..e29bc93c9 100755 --- a/client/src/main/java/org/asynchttpclient/netty/EagerResponseBodyPart.java +++ b/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); + } } diff --git a/client/src/main/java/org/asynchttpclient/netty/LazyResponseBodyPart.java b/client/src/main/java/org/asynchttpclient/netty/LazyResponseBodyPart.java index 3732669d9..59f08dd52 100755 --- a/client/src/main/java/org/asynchttpclient/netty/LazyResponseBodyPart.java +++ b/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; } diff --git a/client/src/main/java/org/asynchttpclient/netty/NettyResponse.java b/client/src/main/java/org/asynchttpclient/netty/NettyResponse.java index 179d796c1..6c9b48649 100755 --- a/client/src/main/java/org/asynchttpclient/netty/NettyResponse.java +++ b/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)); diff --git a/client/src/test/java/org/asynchttpclient/netty/NettyAsyncResponseTest.java b/client/src/test/java/org/asynchttpclient/netty/NettyAsyncResponseTest.java index 30eded248..caadf19b2 100644 --- a/client/src/test/java/org/asynchttpclient/netty/NettyAsyncResponseTest.java +++ b/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 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(); + } }