[azure-core] Fix partial InputStream response buffering

sdk/core/azure-core/CHANGELOG.md

@@ -11,6 +11,8 @@
 - Fixed `AccessTokenCache` and `SimpleTokenCache` leaving concurrent callers waiting forever when the caller that was
   refreshing the token was cancelled, for example by a timeout. Waiting callers now retry the token acquisition when
   the refresh they waited on ends without a token, including when `SimpleTokenCache`'s token supplier completes empty.
+- Fixed streaming of `InputStream`-backed response bodies to emit partial reads without waiting for the response to
+  complete.
 
 ### Other Changes
 

sdk/core/azure-core/src/main/java/com/azure/core/util/FluxUtil.java

@@ -243,7 +243,7 @@ public static Flux<ByteBuffer> createRetriableDownloadFlux(Supplier<Flux<ByteBuf
     }
 
     /**
-     * Converts an {@link InputStream} into a {@link Flux} of {@link ByteBuffer} using a chunk size of 4096.
+     * Converts an {@link InputStream} into a {@link Flux} of {@link ByteBuffer} using a maximum chunk size of 4096.
      * <p>
      * Given that {@link InputStream} is not guaranteed to be replayable the returned {@link Flux} should be considered
      * non-replayable as well.
@@ -266,7 +266,8 @@ public static Flux<ByteBuffer> toFluxByteBuffer(InputStream inputStream) {
      * If the passed {@link InputStream} is {@code null} {@link Flux#empty()} will be returned.
      *
      * @param inputStream The {@link InputStream} to convert into a {@link Flux}.
-     * @param chunkSize The requested size for each {@link ByteBuffer}.
+     * @param chunkSize The maximum size, in bytes, of each {@link ByteBuffer}. Buffers may be smaller when a read returns
+     * fewer bytes and no additional bytes are reported as immediately available.
      * @return A {@link Flux} of {@link ByteBuffer ByteBuffers} that contains the contents of the stream.
      * @throws IllegalArgumentException If {@code chunkSize} is less than or equal to {@code 0}.
      */
@@ -323,25 +324,33 @@ public static Flux<ByteBuffer> toFluxByteBuffer(InputStream inputStream, int chu
             byte[] buffer = new byte[chunkSize];
 
             try {
-                int offset = 0;
+                int readCount;
+                do {
+                    readCount = stream.read(buffer, 0, chunkSize);
+                } while (readCount == 0);
+
+                if (readCount == -1) {
+                    sink.complete();
+                    return stream;
+                }
 
+                int offset = readCount;
                 while (offset < chunkSize) {
-                    int readCount = inputStream.read(buffer, offset, chunkSize - offset);
+                    // Availability is only a batching hint; zero ends coalescing, not the stream.
+                    int available = stream.available();
+                    if (available <= 0) {
+                        break;
+                    }
 
-                    // We have finished reading the stream, trigger onComplete.
-                    if (readCount == -1) {
-                        // If there were bytes read before reaching the end emit the buffer before completing.
-                        if (offset > 0) {
-                            sink.next(ByteBuffer.wrap(buffer, 0, offset));
-                        }
-                        sink.complete();
-                        return stream;
+                    readCount = stream.read(buffer, offset, Math.min(available, chunkSize - offset));
+                    if (readCount <= 0) {
+                        break;
                     }
 
                     offset += readCount;
                 }
 
-                sink.next(ByteBuffer.wrap(buffer));
+                sink.next(ByteBuffer.wrap(buffer, 0, offset));
             } catch (IOException ex) {
                 sink.error(ex);
             }

sdk/core/azure-core/src/test/java/com/azure/core/util/FluxUtilTest.java

@@ -20,8 +20,10 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
 import java.io.ByteArrayInputStream;
@@ -32,6 +34,8 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.PipedInputStream;
+import java.io.PipedOutputStream;
 import java.io.UncheckedIOException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Type;
@@ -474,6 +478,70 @@ private static Stream<Arguments> toFluxByteBufferSupplier() {
             Arguments.arguments(new ByteArrayInputStream(multipleReads), 5432, multipleReads));
     }
 
+    @Test
+    public void toFluxByteBufferWithDefaultAvailable() throws IOException {
+        byte[] expected = "ABCDEFGHI".getBytes(StandardCharsets.UTF_8);
+
+        try (InputStream inputStream = withDefaultAvailable(new ByteArrayInputStream(expected), 3)) {
+            assertEquals(0, inputStream.available());
+
+            StepVerifier.create(FluxUtil.toFluxByteBuffer(inputStream, 8), 0)
+                .thenRequest(1)
+                .assertNext(buffer -> assertEquals("ABC", StandardCharsets.UTF_8.decode(buffer).toString()))
+                .thenRequest(1)
+                .assertNext(buffer -> assertEquals("DEF", StandardCharsets.UTF_8.decode(buffer).toString()))
+                .thenRequest(1)
+                .assertNext(buffer -> assertEquals("GHI", StandardCharsets.UTF_8.decode(buffer).toString()))
+                .thenRequest(1)
+                .expectComplete()
+                .verify(Duration.ofSeconds(2));
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = { false, true })
+    public void toFluxByteBufferEmitsPartialReadBeforeEndOfStream(boolean defaultAvailable) throws IOException {
+        byte[] expected = "data: hello\n\n".getBytes(StandardCharsets.UTF_8);
+
+        try (PipedInputStream inputStream = new PipedInputStream();
+            PipedOutputStream outputStream = new PipedOutputStream(inputStream)) {
+            InputStream stream = defaultAvailable ? withDefaultAvailable(inputStream, expected.length) : inputStream;
+            Flux<ByteBuffer> conversionFlux
+                = FluxUtil.toFluxByteBuffer(stream).subscribeOn(Schedulers.boundedElastic());
+
+            StepVerifier.create(conversionFlux).then(() -> {
+                try {
+                    outputStream.write(expected);
+                    outputStream.flush();
+                } catch (IOException e) {
+                    throw new UncheckedIOException(e);
+                }
+            })
+                .assertNext(buffer -> assertArraysEqual(expected, FluxUtil.byteBufferToArray(buffer)))
+                .thenCancel()
+                .verify(Duration.ofSeconds(2));
+        }
+    }
+
+    private static InputStream withDefaultAvailable(InputStream inputStream, int maxReadSize) {
+        return new InputStream() {
+            @Override
+            public int read() throws IOException {
+                return inputStream.read();
+            }
+
+            @Override
+            public int read(byte[] buffer, int offset, int length) throws IOException {
+                return inputStream.read(buffer, offset, Math.min(length, maxReadSize));
+            }
+
+            @Override
+            public void close() throws IOException {
+                inputStream.close();
+            }
+        };
+    }
+
     @Test
     public void toFluxByteBufferMultipleSubscriptions() {
         byte[] singleRead = new byte[4096];

sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTests.java

@@ -99,6 +99,7 @@
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -114,6 +115,8 @@
 import static com.azure.core.validation.http.HttpClientTestsServer.HUGE_HEADER_VALUE;
 import static com.azure.core.validation.http.HttpClientTestsServer.INVALID_HEADER_RESPONSE;
 import static com.azure.core.validation.http.HttpClientTestsServer.PLAIN_RESPONSE;
+import static com.azure.core.validation.http.HttpClientTestsServer.PARTIAL_RESPONSE;
+import static com.azure.core.validation.http.HttpClientTestsServer.PARTIAL_RESPONSE_BYTES;
 import static com.azure.core.validation.http.HttpClientTestsServer.RETURN_BYTES;
 import static com.azure.core.validation.http.HttpClientTestsServer.UTF_16BE_BOM_RESPONSE;
 import static com.azure.core.validation.http.HttpClientTestsServer.UTF_16LE_BOM_RESPONSE;
@@ -131,6 +134,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -203,6 +207,23 @@ public void plainResponse() {
         assertEquals(expected, actual);
     }
 
+    /**
+     * Tests that a partial response is emitted without waiting for the response to end.
+     */
+    @SyncAsyncTest
+    public void partialResponseEmittedBeforeEndOfStream() {
+        byte[] actual = SyncAsyncExtension.execute(() -> assertTimeoutPreemptively(Duration.ofSeconds(3), () -> {
+            HttpResponse response = createHttpClient()
+                .sendSync(new HttpRequest(HttpMethod.GET, getRequestUrl(PARTIAL_RESPONSE)), Context.NONE);
+            return collectPartialResponse(response).block();
+        }), () -> createHttpClient().send(new HttpRequest(HttpMethod.GET, getRequestUrl(PARTIAL_RESPONSE)))
+            .flatMap(this::collectPartialResponse)
+            .timeout(Duration.ofSeconds(3)));
+
+        assertNotNull(actual);
+        assertArraysEqual(PARTIAL_RESPONSE_BYTES, actual);
+    }
+
     /**
      * Tests that a response with a 'Content-Type' header encodes using the specified charset.
      */
@@ -815,6 +836,14 @@ private Mono<String> sendRequest(String requestPath) {
             .flatMap(HttpResponse::getBodyAsString);
     }
 
+    private Mono<byte[]> collectPartialResponse(HttpResponse response) {
+        AtomicInteger received = new AtomicInteger();
+        Flux<ByteBuffer> responseBody = response.getBody()
+            .takeUntil(buffer -> received.addAndGet(buffer.remaining()) >= PARTIAL_RESPONSE_BYTES.length);
+
+        return FluxUtil.collectBytesInByteBufferStream(responseBody).doFinally(ignored -> response.close());
+    }
+
     private String sendRequestSync(String requestPath) {
         try (HttpResponse httpResponse
             = createHttpClient().sendSync(new HttpRequest(HttpMethod.GET, getRequestUrl(requestPath)), Context.NONE)) {

sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTestsServer.java

@@ -36,6 +36,7 @@
 @Execution(ExecutionMode.SAME_THREAD)
 public final class HttpClientTestsServer {
     static final String PLAIN_RESPONSE = "plainBytesNoHeader";
+    static final String PARTIAL_RESPONSE = "partialResponse";
     static final String HEADER_RESPONSE = "plainBytesWithHeader";
     static final String INVALID_HEADER_RESPONSE = "plainBytesInvalidHeader";
     static final String UTF_8_BOM_RESPONSE = "utf8BomBytes";
@@ -56,6 +57,7 @@ public final class HttpClientTestsServer {
 
     private static final String HELLO_WORLD = "Hello World!";
     static final byte[] RETURN_BYTES = HELLO_WORLD.getBytes(StandardCharsets.UTF_8);
+    static final byte[] PARTIAL_RESPONSE_BYTES = "data: hello\n\n".getBytes(StandardCharsets.UTF_8);
     static final HttpHeaderName HUGE_HEADER_NAME = HttpHeaderName.fromString("x-huge-header");
     static final String HUGE_HEADER_VALUE;
 
@@ -119,6 +121,17 @@ public static LocalTestServer getHttpClientTestsServer() {
                 resp.getHttpOutput().complete(Callback.NOOP);
             } else if (get && pathMatches(path, PLAIN_RESPONSE)) {
                 handleRequest(resp, "application/octet-stream", RETURN_BYTES);
+            } else if (get && pathMatches(path, PARTIAL_RESPONSE)) {
+                resp.setContentType("text/event-stream");
+                resp.getHttpOutput().write(PARTIAL_RESPONSE_BYTES);
+                resp.getHttpOutput().flush();
+                try {
+                    Thread.sleep(5000);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new ServletException(e);
+                }
+                resp.getHttpOutput().complete(Callback.NOOP);
             } else if (get && pathMatches(path, HEADER_RESPONSE)) {
                 handleRequest(resp, "charset=UTF-16BE", RETURN_BYTES);
             } else if (get && pathMatches(path, INVALID_HEADER_RESPONSE)) {