Update BufferedHttpResponse (Azure#24128)

Update how BufferedHttpResponse Buffers Response Content

Author: alzimmermsft

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/BufferedHttpResponse.java

@@ -7,18 +7,26 @@
 import com.azure.core.http.HttpResponse;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.FluxUtil;
+import com.azure.core.util.logging.ClientLogger;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static com.azure.core.util.FluxUtil.monoError;
 
 /**
  * HTTP response which will buffer the response's body when/if it is read.
  */
 public final class BufferedHttpResponse extends HttpResponse {
+    private final ClientLogger logger = new ClientLogger(BufferedHttpResponse.class);
+
     private final HttpResponse innerHttpResponse;
-    private final Flux<ByteBuffer> cachedBody;
+    private final Mono<List<ByteBuffer>> cachedBody;
+    private final AtomicLong cachedBodySize = new AtomicLong();
 
     /**
      * Creates a buffered HTTP response.
@@ -28,12 +36,13 @@ public final class BufferedHttpResponse extends HttpResponse {
     public BufferedHttpResponse(HttpResponse innerHttpResponse) {
         super(innerHttpResponse.getRequest());
         this.innerHttpResponse = innerHttpResponse;
-        this.cachedBody = FluxUtil.collectBytesFromNetworkResponse(innerHttpResponse.getBody(),
-            innerHttpResponse.getHeaders())
-            .map(ByteBuffer::wrap)
-            .flux()
-            .cache()
-            .map(ByteBuffer::duplicate);
+        this.cachedBody = innerHttpResponse.getBody()
+            .map(buffer -> {
+                cachedBodySize.addAndGet(buffer.remaining());
+                return ByteBuffer.wrap(FluxUtil.byteBufferToArray(buffer));
+            })
+            .collectList()
+            .cache();
     }
 
     @Override
@@ -53,23 +62,35 @@ public HttpHeaders getHeaders() {
 
     @Override
     public Flux<ByteBuffer> getBody() {
-        return cachedBody;
+        return cachedBody.flatMapMany(Flux::fromIterable).map(ByteBuffer::duplicate);
     }
 
     @Override
     public Mono<byte[]> getBodyAsByteArray() {
-        return cachedBody.next().map(ByteBuffer::array);
+        // Check that the body would fit into a byte array before spending time to create the merged byte array.
+        return (cachedBodySize.get() > Integer.MAX_VALUE)
+            ? monoError(logger, new IllegalStateException(
+                "Response with body size " + cachedBodySize.get() + " doesn't fit into a byte array."))
+            : FluxUtil.collectBytesInByteBufferStream(getBody(), (int) cachedBodySize.get());
     }
 
     @Override
     public Mono<String> getBodyAsString() {
-        return getBodyAsByteArray().map(bytes ->
-            CoreUtils.bomAwareToString(bytes, innerHttpResponse.getHeaderValue("Content-Type")));
+        // Check that the body would fit into a String before spending the time to create the String.
+        return (cachedBodySize.get() > Integer.MAX_VALUE)
+            ? monoError(logger, new IllegalStateException(
+                "Response with body size " + cachedBodySize.get() + " doesn't fit into a String."))
+            : getBodyAsByteArray().map(bytes ->
+                CoreUtils.bomAwareToString(bytes, innerHttpResponse.getHeaderValue("Content-Type")));
     }
 
     @Override
     public Mono<String> getBodyAsString(Charset charset) {
-        return getBodyAsByteArray().map(bytes -> new String(bytes, charset));
+        // Check that the body would fit into a String before spending the time to create the String.
+        return (cachedBodySize.get() > Integer.MAX_VALUE)
+            ? monoError(logger, new IllegalStateException(
+                "Response with body size " + cachedBodySize.get() + " doesn't fit into a String."))
+            : getBodyAsByteArray().map(bytes -> new String(bytes, charset));
     }
 
     @Override