diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md index 5af4c8c55b48..19870b9b3bef 100644 --- a/sdk/core/azure-core/CHANGELOG.md +++ b/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 diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/util/FluxUtil.java b/sdk/core/azure-core/src/main/java/com/azure/core/util/FluxUtil.java index 1356122138cb..e79d45a168d2 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/util/FluxUtil.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/util/FluxUtil.java @@ -243,7 +243,7 @@ public static Flux createRetriableDownloadFlux(Supplier * 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 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 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); } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/util/FluxUtilTest.java b/sdk/core/azure-core/src/test/java/com/azure/core/util/FluxUtilTest.java index fba14c833fd8..1ae1be909d90 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/util/FluxUtilTest.java +++ b/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 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 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]; diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTests.java index 541a9c69e82a..a1fd34e85d77 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTests.java +++ b/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 sendRequest(String requestPath) { .flatMap(HttpResponse::getBodyAsString); } + private Mono collectPartialResponse(HttpResponse response) { + AtomicInteger received = new AtomicInteger(); + Flux 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)) { diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTestsServer.java b/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTestsServer.java index bd20025b8e89..817adf55eec3 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/validation/http/HttpClientTestsServer.java +++ b/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)) {