From b93ca12cfa7470c22a52d82b3e3eb5fb3406c365 Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Tue, 18 Aug 2026 23:32:15 +0800 Subject: [PATCH 1/8] Preserve SSE response streams in RestProxy Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daa3657a-91d5-48fa-b7ae-6983091cab6a --- .../core/http/policy/HttpLoggingPolicy.java | 6 +- .../http/rest/AsyncRestProxy.java | 48 ++++++---- .../http/rest/RestProxyBase.java | 35 ++++++- .../http/rest/SyncRestProxy.java | 18 +++- .../core/implementation/util/HttpUtils.java | 91 +++++++++++++++++++ .../http/policy/HttpLoggingPolicyTests.java | 85 +++++++++++++++++ .../http/rest/AsyncRestProxyTests.java | 48 ++++++++++ .../http/rest/SyncRestProxyTests.java | 48 ++++++++++ .../implementation/util/HttpUtilsTests.java | 31 +++++++ 9 files changed, 385 insertions(+), 25 deletions(-) create mode 100644 sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java index 20e2c2f7c9d9..6a4e3f172c51 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java @@ -21,6 +21,7 @@ import com.azure.core.implementation.util.ByteArrayContent; import com.azure.core.implementation.util.ByteBufferContent; import com.azure.core.implementation.util.HttpHeadersAccessHelper; +import com.azure.core.implementation.util.HttpUtils; import com.azure.core.implementation.util.InputStreamContent; import com.azure.core.implementation.util.SerializableContent; import com.azure.core.implementation.util.StringContent; @@ -506,8 +507,8 @@ private Long getAndLogContentLength(HttpHeaders headers, LoggingEventBuilder log /* * Determines if the request or response body should be logged. * - *

The request or response body is logged if the Content-Type is not "application/octet-stream" and the body - * isn't empty and is less than 16KB in size.

+ *

The request or response body is logged if the Content-Type is neither "application/octet-stream" nor + * "text/event-stream" and the body isn't empty and is less than 16KB in size.

* * @param contentTypeHeader Content-Type header value. * @@ -518,6 +519,7 @@ private Long getAndLogContentLength(HttpHeaders headers, LoggingEventBuilder log private static boolean shouldBodyBeLogged(String contentTypeHeader, Long contentLength) { return contentLength != null && !ContentType.APPLICATION_OCTET_STREAM.equalsIgnoreCase(contentTypeHeader) + && !HttpUtils.isTextEventStreamContentType(contentTypeHeader) && contentLength != 0 && contentLength < MAX_BODY_LOG_SIZE; } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java index 2548bec2a64f..1bf161a9431f 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java @@ -13,6 +13,7 @@ import com.azure.core.http.rest.StreamResponse; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.serializer.HttpResponseDecoder; +import com.azure.core.implementation.util.HttpUtils; import com.azure.core.util.Base64Url; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; @@ -46,8 +47,6 @@ */ public class AsyncRestProxy extends RestProxyBase { - private static final String TEXT_EVENT_STREAM = "text/event-stream"; - /** * Create a RestProxy. * @@ -83,6 +82,8 @@ public Object invoke(Object proxy, Method method, RequestOptions options, EnumSe requestCallback.accept(request); } + context = updateRequestContext(request, context); + final Context finalContext = context; final Mono asyncResponse = RestProxyUtils.validateLengthAsync(request).flatMap(r -> { // correlates logs @@ -142,7 +143,10 @@ private Mono ensureExpectedStatus( } private Mono handleRestResponseReturnType(final HttpResponseDecoder.HttpDecodedResponse response, - final SwaggerMethodParser methodParser, final Type entityType) { + final SwaggerMethodParser methodParser, final Type entityType, boolean preserveResponseBodyAsStream) { + final boolean shouldPreserveResponseBodyAsStream = preserveResponseBodyAsStream + || HttpUtils.isTextEventStreamContentType( + response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); if (methodParser.isStreamResponse()) { return Mono.fromSupplier(() -> new StreamResponse(response.getSourceResponse())); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) { @@ -151,15 +155,20 @@ private Mono handleRestResponseReturnType(final HttpResponseDecoder.HttpDecod return response.getSourceResponse() .getBody() .ignoreElements() - .then(Mono.fromCallable(() -> createResponse(response, entityType, null))); + .then(Mono.fromCallable( + () -> createResponse(response, entityType, null, shouldPreserveResponseBodyAsStream))); } else { - return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType) - .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject)) - .switchIfEmpty(Mono.fromCallable(() -> createResponse(response, entityType, null))); + return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType, + shouldPreserveResponseBodyAsStream) + .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject, + shouldPreserveResponseBodyAsStream)) + .switchIfEmpty(Mono.fromCallable( + () -> createResponse(response, entityType, null, shouldPreserveResponseBodyAsStream))); } } else { // For now, we're just throwing if the Maybe didn't emit a value. - return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType); + return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType, + shouldPreserveResponseBodyAsStream); } } @@ -177,7 +186,7 @@ private static Function> decodeBytes(HttpResponseDecoder.Ht } static Mono handleBodyReturnType(HttpResponse sourceResponse, Function> getDecodedBody, - SwaggerMethodParser methodParser, Type entityType) { + SwaggerMethodParser methodParser, Type entityType, boolean responseBodyStreaming) { final int responseStatusCode = sourceResponse.getStatusCode(); final HttpMethod httpMethod = methodParser.getHttpMethod(); final Type returnValueWireType = methodParser.getReturnValueWireType(); @@ -207,9 +216,9 @@ static Mono handleBodyReturnType(HttpResponse sourceResponse, Function handleBodyReturnType(HttpResponse sourceResponse, Function handleBodyReturnType(HttpResponse sourceResponse, Function> getDecodedBody, + SwaggerMethodParser methodParser, Type entityType) { + return handleBodyReturnType(sourceResponse, getDecodedBody, methodParser, entityType, false); + } + /** * Handle the provided asynchronous HTTP response and return the deserialized value. * @@ -238,6 +252,7 @@ private Object handleRestReturnType(Mono errorOptionsSet) { final Mono asyncExpectedResponse = endSpanWhenDone( ensureExpectedStatus(asyncHttpDecodedResponse, methodParser, options, errorOptionsSet), context); + final boolean preserveResponseBodyAsStream = HttpUtils.shouldPreserveResponseBodyAsStream(context); final Object result; if (TypeUtil.isTypeOrSubTypeOf(returnType, Mono.class)) { @@ -247,8 +262,8 @@ private Object handleRestReturnType(Mono> - result = asyncExpectedResponse - .flatMap(response -> handleRestResponseReturnType(response, methodParser, monoTypeParam)); + result = asyncExpectedResponse.flatMap(response -> handleRestResponseReturnType(response, methodParser, + monoTypeParam, preserveResponseBodyAsStream)); } } else if (FluxUtil.isFluxByteBuffer(returnType)) { // ProxyMethod ReturnType: Flux @@ -261,9 +276,8 @@ private Object handleRestReturnType(Mono handleRestResponseReturnType(httpResponse, methodParser, returnType)) - .block(); + result = asyncExpectedResponse.flatMap(httpResponse -> handleRestResponseReturnType(httpResponse, + methodParser, returnType, preserveResponseBodyAsStream)).block(); } return result; } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java index a33abec128a8..2c70a09b1106 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java @@ -22,6 +22,7 @@ import com.azure.core.http.rest.RequestOptions; import com.azure.core.http.rest.Response; import com.azure.core.http.rest.ResponseBase; +import com.azure.core.http.rest.SimpleResponse; import com.azure.core.implementation.ReflectiveInvoker; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.http.UnexpectedExceptionInformation; @@ -36,12 +37,14 @@ import com.azure.core.util.tracing.Tracer; import reactor.core.Exceptions; +import java.io.Closeable; import java.io.IOException; import java.lang.reflect.Method; import java.lang.reflect.Type; import java.net.URL; import java.nio.charset.StandardCharsets; import java.util.EnumSet; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import static com.azure.core.util.FluxUtil.monoError; @@ -149,6 +152,14 @@ protected abstract Object invoke(Object proxy, Method method, RequestOptions opt EnumSet errorOptions, Consumer httpRequestConsumer, SwaggerMethodParser methodParser, HttpRequest request, Context context); + final Context updateRequestContext(HttpRequest request, Context context) { + if (HttpUtils.acceptsTextEventStream(request.getHeaders().getValue(HttpHeaderName.ACCEPT))) { + return context.addData(HttpUtils.AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM, true); + } + + return context; + } + /** * Update the request with the provided configuration. * @@ -165,13 +176,14 @@ public abstract void updateRequest(RequestDataConfiguration requestDataConfigura * @param response the decoded response * @param entityType the type of the response entity * @param bodyAsObject the response body as an object + * @param responseBodyStreaming whether the response owns an unconsumed streaming body * @return the {@link Response} * @throws RuntimeException If the response type is a PagedResponse and the bodyAsObject is not an instance of * Page. */ @SuppressWarnings({ "unchecked", "rawtypes" }) public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, Type entityType, - Object bodyAsObject) { + Object bodyAsObject, boolean responseBodyStreaming) { final Class> cls = (Class>) TypeUtil.getRawClass(entityType); final HttpResponse httpResponse = response.getSourceResponse(); @@ -185,6 +197,10 @@ public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, // If the type is either the Response or PagedResponse interface from azure-core a new instance of either // ResponseBase or PagedResponseBase can be returned. if (cls.equals(Response.class)) { + if (responseBodyStreaming) { + return cls.cast(new CloseableResponse<>(httpResponse, bodyAsObject)); + } + // For Response return a new instance of ResponseBase cast to the class. return cls.cast(new ResponseBase<>(request, statusCode, headers, bodyAsObject, decodedHeaders)); } else if (cls.equals(PagedResponse.class)) { @@ -213,6 +229,23 @@ public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, return RESPONSE_CONSTRUCTORS_CACHE.invoke(constructorReflectiveInvoker, response, bodyAsObject); } + private static final class CloseableResponse extends SimpleResponse implements Closeable { + private final AtomicBoolean closed = new AtomicBoolean(); + private final HttpResponse response; + + private CloseableResponse(HttpResponse response, T value) { + super(response.getRequest(), response.getStatusCode(), response.getHeaders(), value); + this.response = response; + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + response.close(); + } + } + } + /** * Starts the tracing span for the current service call, additionally set metadata attributes on the span by passing * additional context information. diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index e2b3ca0c95c8..148f34820634 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -3,6 +3,7 @@ package com.azure.core.implementation.http.rest; +import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpMethod; import com.azure.core.http.HttpPipeline; import com.azure.core.http.HttpRequest; @@ -13,6 +14,7 @@ import com.azure.core.implementation.ImplUtils; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.serializer.HttpResponseDecoder; +import com.azure.core.implementation.util.HttpUtils; import com.azure.core.util.Base64Url; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; @@ -76,6 +78,8 @@ public Object invoke(Object proxy, Method method, RequestOptions options, EnumSe requestCallback.accept(request); } + context = updateRequestContext(request, context); + if (request.getBodyAsBinaryData() != null) { request.setBody(RestProxyUtils.validateLengthSync(request)); } @@ -136,19 +140,22 @@ private HttpResponseDecoder.HttpDecodedResponse ensureExpectedStatus( } private Object handleRestResponseReturnType(HttpResponseDecoder.HttpDecodedResponse response, - SwaggerMethodParser methodParser, Type entityType) { + SwaggerMethodParser methodParser, Type entityType, boolean responseBodyStreaming) { + responseBodyStreaming = responseBodyStreaming + || HttpUtils.isTextEventStreamContentType( + response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); if (methodParser.isStreamResponse()) { return new StreamResponse(response.getSourceResponse()); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) { final Type bodyType = TypeUtil.getRestResponseBodyType(entityType); if (TypeUtil.isTypeOrSubTypeOf(bodyType, Void.class)) { response.getSourceResponse().close(); - return createResponse(response, entityType, null); + return createResponse(response, entityType, null, responseBodyStreaming); } else { Object bodyAsObject = handleBodyReturnType(response, methodParser, bodyType); - Response httpResponse = createResponse(response, entityType, bodyAsObject); + Response httpResponse = createResponse(response, entityType, bodyAsObject, responseBodyStreaming); if (httpResponse == null) { - return createResponse(response, entityType, null); + return createResponse(response, entityType, null, responseBodyStreaming); } return httpResponse; } @@ -218,7 +225,8 @@ private Object handleRestReturnType(HttpResponseDecoder.HttpDecodedResponse http } else { // ProxyMethod ReturnType: T where T != async (Mono, Flux) or sync Void // Block the deserialization until a value T is received - result = handleRestResponseReturnType(httpDecodedResponse, methodParser, returnType); + result = handleRestResponseReturnType(httpDecodedResponse, methodParser, returnType, + HttpUtils.shouldPreserveResponseBodyAsStream(context)); } return result; } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java index 0fd02b8f9f3a..bf9f41a7512e 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java @@ -3,9 +3,13 @@ package com.azure.core.implementation.util; import com.azure.core.util.Configuration; +import com.azure.core.util.Context; import com.azure.core.util.logging.ClientLogger; import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import static com.azure.core.util.Configuration.PROPERTY_AZURE_REQUEST_CONNECT_TIMEOUT; import static com.azure.core.util.Configuration.PROPERTY_AZURE_REQUEST_READ_TIMEOUT; @@ -17,6 +21,7 @@ * Utilities shared with HttpClient implementations. */ public final class HttpUtils { + private static final String TEXT_EVENT_STREAM = "text/event-stream"; private static final ClientLogger LOGGER = new ClientLogger(HttpUtils.class); private static final Duration MINIMUM_TIMEOUT = Duration.ofMillis(1); @@ -43,6 +48,12 @@ public final class HttpUtils { */ public static final String AZURE_EAGERLY_READ_RESPONSE = "azure-eagerly-read-response"; + /** + * Context key that instructs REST proxy response ownership and decoding to preserve the response body as a live + * stream. HTTP client implementations do not consume this key. + */ + public static final String AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM = "azure-preserve-response-body-as-stream"; + /** * Context key used to indicate to an HttpClient implementation if the response body should be ignored and eagerly * drained from the network. @@ -60,6 +71,86 @@ public final class HttpUtils { */ public static final String AZURE_EAGERLY_CONVERT_HEADERS = "azure-eagerly-convert-headers"; + /** + * Determines whether the response body must be preserved as a live stream. + * + * @param context Contextual information about the request. + * @return Whether the response body must be preserved as a live stream. + */ + public static boolean shouldPreserveResponseBodyAsStream(Context context) { + return Boolean.TRUE.equals(context.getData(AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM).orElse(false)); + } + + /** + * Determines whether an Accept header contains a {@code text/event-stream} media range. + * + * @param headerValue The header value. + * @return Whether the header contains a {@code text/event-stream} media range. + */ + public static boolean acceptsTextEventStream(String headerValue) { + if (headerValue == null) { + return false; + } + + for (String value : splitHeaderValue(headerValue, ',')) { + if (isTextEventStreamContentType(value)) { + return true; + } + } + + return false; + } + + /** + * Determines whether a Content-Type header identifies exactly one {@code text/event-stream} representation. + * Charset parameters don't affect this determination as event streams are always decoded as UTF-8. + * + * @param headerValue The Content-Type header value. + * @return Whether the header identifies a {@code text/event-stream} representation. + */ + public static boolean isTextEventStreamContentType(String headerValue) { + if (headerValue == null) { + return false; + } + + List mediaTypeAndParameters = splitHeaderValue(headerValue, ';'); + if (mediaTypeAndParameters.size() == 0 + || !TEXT_EVENT_STREAM.equalsIgnoreCase(mediaTypeAndParameters.get(0).trim()) + || splitHeaderValue(headerValue, ',').size() != 1) { + return false; + } + + return true; + } + + private static List splitHeaderValue(String value, char delimiter) { + List segments = new ArrayList<>(); + int start = 0; + boolean quoted = false; + boolean escaped = false; + + for (int i = 0; i < value.length(); i++) { + char character = value.charAt(i); + if (escaped) { + escaped = false; + } else if (quoted && character == '\\') { + escaped = true; + } else if (character == '"') { + quoted = !quoted; + } else if (!quoted && character == delimiter) { + segments.add(value.substring(start, i)); + start = i + 1; + } + } + + if (quoted) { + return Collections.emptyList(); + } + + segments.add(value.substring(start)); + return segments; + } + /** * Gets the default connect timeout. * diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java index e21f13423dcb..3db53cbc3b66 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java @@ -16,6 +16,7 @@ import com.azure.core.implementation.AccessibleByteArrayOutputStream; import com.azure.core.implementation.accesshelpers.ClientLoggerAccessHelper; import com.azure.core.implementation.logging.DefaultLogger; +import com.azure.core.implementation.util.HttpUtils; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.core.util.FluxUtil; @@ -28,6 +29,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.parallel.Execution; import org.junit.jupiter.api.parallel.ExecutionMode; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; @@ -58,6 +60,7 @@ import static com.azure.core.CoreTestUtils.createUrl; import static com.azure.core.http.HttpHeaderName.X_MS_REQUEST_ID; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -192,6 +195,32 @@ public void validateLoggingDoesNotConsumeRequestSync(BinaryData requestBody, byt expectedRequest.assertEqual(messages.get(0), logOptions, LogLevel.INFORMATIONAL); } + @Test + public void textEventStreamRequestBodiesAreNotLogged() { + byte[] data = "request event data".getBytes(StandardCharsets.UTF_8); + AtomicInteger requestCount = new AtomicInteger(); + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(new HttpLoggingPolicy(new HttpLogOptions().setLogLevel(HttpLogDetailLevel.BODY))) + .httpClient(request -> FluxUtil.collectBytesInByteBufferStream(request.getBody()).doOnSuccess(bytes -> { + assertArraysEqual(data, bytes); + requestCount.incrementAndGet(); + }).then(Mono.empty())) + .build(); + + HttpRequest asyncRequest = new HttpRequest(HttpMethod.POST, "https://test.com/async") + .setHeader(HttpHeaderName.CONTENT_TYPE, "Text/Event-Stream; charset=utf-8") + .setBody(BinaryData.fromBytes(data)); + StepVerifier.create(pipeline.send(asyncRequest)).verifyComplete(); + + HttpRequest syncRequest = new HttpRequest(HttpMethod.POST, "https://test.com/sync") + .setHeader(HttpHeaderName.CONTENT_TYPE, "Text/Event-Stream; charset=utf-8") + .setBody(BinaryData.fromBytes(data)); + pipeline.sendSync(syncRequest, Context.NONE); + + assertEquals(2, requestCount.get()); + assertFalse(convertOutputStreamToString(logCaptureStream).contains(new String(data, StandardCharsets.UTF_8))); + } + /** * Tests that logging the response body doesn't consume the stream before it is returned from the service call. */ @@ -244,6 +273,46 @@ public void validateLoggingDoesNotConsumeResponseSync(BinaryData responseBody, b assertTrue(logString.contains(new String(data, StandardCharsets.UTF_8))); } + @ParameterizedTest(name = "[{index}] {displayName}") + @MethodSource("responseLoggingSupplier") + public void responseLoggingUsesActualContentType(String contentType, boolean useStreamingContext, + int expectedBufferCount, boolean expectBodyLogged) { + byte[] data = "streaming response".getBytes(StandardCharsets.UTF_8); + AtomicInteger bufferCount = new AtomicInteger(); + HttpRequest request = new HttpRequest(HttpMethod.GET, "https://test.com/responseLoggingUsesActualContentType"); + HttpHeaders responseHeaders = new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, contentType) + .set(HttpHeaderName.CONTENT_LENGTH, Integer.toString(data.length)); + + HttpPipeline pipeline = new HttpPipelineBuilder() + .policies(new HttpLoggingPolicy(new HttpLogOptions().setLogLevel(HttpLogDetailLevel.BODY))) + .httpClient(ignored -> Mono.just(new BufferTrackingHttpResponse(ignored, responseHeaders, + Flux.just(ByteBuffer.wrap(data)), bufferCount))) + .build(); + + Context context = getCallerMethodContext("streamingResponsesAreNotBuffered", LogLevel.INFORMATIONAL); + if (useStreamingContext) { + context = context.addData(HttpUtils.AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM, true); + } + + try (HttpResponse response = pipeline.send(request, context).block()) { + assertNotNull(response); + assertArraysEqual(data, response.getBodyAsBinaryData().toBytes()); + } + + try (HttpResponse response = pipeline.sendSync(request, context)) { + assertArraysEqual(data, response.getBodyAsBinaryData().toBytes()); + } + + assertEquals(expectedBufferCount, bufferCount.get()); + assertEquals(expectBodyLogged, + convertOutputStreamToString(logCaptureStream).contains(new String(data, StandardCharsets.UTF_8))); + } + + private static Stream responseLoggingSupplier() { + return Stream.of(Arguments.of("Text/Event-Stream; charset=utf-8", false, 0, false), + Arguments.of(ContentType.APPLICATION_JSON, true, 2, true)); + } + private static Stream validateLoggingDoesNotConsumeSupplierSync() { byte[] data = "this is a test".getBytes(StandardCharsets.UTF_8); @@ -351,6 +420,22 @@ public Mono getBodyAsString(Charset charset) { } } + private static final class BufferTrackingHttpResponse extends MockHttpResponse { + private final AtomicInteger bufferCount; + + private BufferTrackingHttpResponse(HttpRequest request, HttpHeaders headers, Flux body, + AtomicInteger bufferCount) { + super(request, headers, body); + this.bufferCount = bufferCount; + } + + @Override + public HttpResponse buffer() { + bufferCount.incrementAndGet(); + return super.buffer(); + } + } + @ParameterizedTest(name = "[{index}] {displayName}") @MethodSource("logOptionsSupplier") public void loggingIncludesRetryCount(HttpLogOptions logOptions) { diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java index af0ddd076d72..de451efe2e1e 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java @@ -3,33 +3,44 @@ package com.azure.core.implementation.http.rest; +import com.azure.core.annotation.ExpectedResponses; import com.azure.core.annotation.Get; import com.azure.core.annotation.Head; import com.azure.core.annotation.Host; import com.azure.core.annotation.ServiceInterface; import com.azure.core.http.ContentType; +import com.azure.core.http.HttpClient; import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpHeaders; +import com.azure.core.http.HttpPipelineBuilder; import com.azure.core.http.HttpResponse; import com.azure.core.http.MockHttpResponse; +import com.azure.core.http.rest.RequestOptions; +import com.azure.core.http.rest.Response; +import com.azure.core.http.rest.RestProxy; import com.azure.core.util.BinaryData; +import com.azure.core.util.Context; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; import reactor.test.StepVerifier; +import java.io.Closeable; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Type; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Stream; import static com.azure.core.CoreTestUtils.assertArraysEqual; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -54,6 +65,10 @@ private interface MockService { @Get("getStreamResponse") Flux getStreamResponse(); + + @Get("getStreamingResponse") + @ExpectedResponses({ 200 }) + Mono> getStreamingResponse(RequestOptions options, Context context); } @BeforeEach @@ -193,4 +208,37 @@ public static Stream getResponseHeaderAndReplayability() { Arguments.of(new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, ContentType.APPLICATION_JSON), true), Arguments.of(new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, "application/xml"), true)); } + + @ParameterizedTest + @MethodSource("streamingResponseOwnershipSupplier") + public void streamingResponseIsUnbufferedAndCloseable(String accept, String contentType) throws IOException { + byte[] expectedBytes = "hello".getBytes(StandardCharsets.UTF_8); + AtomicBoolean responseClosed = new AtomicBoolean(); + HttpClient client = request -> Mono.just(new MockHttpResponse(request, 200, + new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, contentType), expectedBytes) { + @Override + public void close() { + responseClosed.set(true); + super.close(); + } + }); + MockService service = RestProxy.create(MockService.class, new HttpPipelineBuilder().httpClient(client).build()); + RequestOptions options = new RequestOptions(); + if (accept != null) { + options.setHeader(HttpHeaderName.ACCEPT, accept); + } + + Response response = service.getStreamingResponse(options, Context.NONE).block(); + + assertTrue(response instanceof Closeable); + assertFalse(response.getValue().isReplayable()); + assertArraysEqual(expectedBytes, response.getValue().toBytes()); + ((Closeable) response).close(); + assertTrue(responseClosed.get()); + } + + private static Stream streamingResponseOwnershipSupplier() { + return Stream.of(Arguments.of("text/event-stream", ContentType.APPLICATION_JSON), + Arguments.of(null, "text/event-stream; charset=utf-8")); + } } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java index 682ec7dc5f7f..99809e7b1e63 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java @@ -13,6 +13,7 @@ import com.azure.core.annotation.ServiceInterface; import com.azure.core.http.HttpClient; import com.azure.core.http.HttpHeaderName; +import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; import com.azure.core.http.HttpPipeline; import com.azure.core.http.HttpPipelineBuilder; @@ -32,6 +33,7 @@ import reactor.core.publisher.Mono; import java.io.ByteArrayInputStream; +import java.io.Closeable; import java.io.IOException; import java.io.InputStream; import java.util.Collections; @@ -67,6 +69,10 @@ Response testMethod(@BodyParam("application/octet-stream") BinaryData data @Put("my/url/path") @ExpectedResponses({ 200 }) Response testInputStreamResponse(Context context); + + @Get("my/url/path") + @ExpectedResponses({ 200 }) + Response getStreamingResponse(RequestOptions options, Context context); } @Test @@ -193,6 +199,48 @@ public void testInputStream() throws IOException { assertEquals("hello", new String(bytes)); } + @ParameterizedTest + @MethodSource("streamingResponseOwnershipSupplier") + public void streamingResponseIsCloseable(String accept, String contentType) throws IOException { + AtomicBoolean responseClosed = new AtomicBoolean(); + HttpClient client = new HttpClient() { + @Override + public Mono send(HttpRequest request) { + return Mono.error(new IllegalStateException("Async Send API was Invoked.")); + } + + @Override + public HttpResponse sendSync(HttpRequest request, Context context) { + return new MockHttpResponse(request, 200, + new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, contentType), "hello".getBytes()) { + @Override + public void close() { + responseClosed.set(true); + super.close(); + } + }; + } + }; + TestInterface service + = RestProxy.create(TestInterface.class, new HttpPipelineBuilder().httpClient(client).build()); + RequestOptions options = new RequestOptions(); + if (accept != null) { + options.setHeader(HttpHeaderName.ACCEPT, accept); + } + + Response response = service.getStreamingResponse(options, Context.NONE); + + assertTrue(response instanceof Closeable); + assertEquals("hello", response.getValue().toString()); + ((Closeable) response).close(); + assertTrue(responseClosed.get()); + } + + private static Stream streamingResponseOwnershipSupplier() { + return Stream.of(Arguments.of("text/event-stream", "application/json"), + Arguments.of(null, "text/event-stream; charset=utf-8")); + } + private static Stream mergeRequestOptionsContextSupplier() { Map twoValuesMap = new HashMap<>(); twoValuesMap.put("key", "value"); diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java new file mode 100644 index 000000000000..82448e7eeb80 --- /dev/null +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.core.implementation.util; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class HttpUtilsTests { + @Test + public void acceptsTextEventStreamIgnoresParameters() { + assertTrue(HttpUtils.acceptsTextEventStream("text/event-stream;q=0")); + assertTrue(HttpUtils.acceptsTextEventStream("application/json, text/event-stream;q=invalid")); + assertFalse(HttpUtils.acceptsTextEventStream("application/json; note=\"text/event-stream, q=1\"")); + assertTrue(HttpUtils.acceptsTextEventStream("text/event-stream; note=\"x,y;q=0.5\"")); + } + + @Test + public void textEventStreamContentTypeRequiresSingleMediaType() { + assertTrue(HttpUtils.isTextEventStreamContentType("Text/Event-Stream; charset=utf-8")); + assertTrue(HttpUtils.isTextEventStreamContentType("text/event-stream; charset=\"UTF-8\"")); + assertTrue(HttpUtils.isTextEventStreamContentType("text/event-stream; charset=iso-8859-1")); + assertTrue(HttpUtils.isTextEventStreamContentType("text/event-stream; charset=utf-16")); + assertTrue(HttpUtils.isTextEventStreamContentType("text/event-stream; charset=not-a-charset; charset=utf-16")); + assertFalse(HttpUtils.isTextEventStreamContentType("application/json, text/event-stream")); + assertTrue(HttpUtils.isTextEventStreamContentType("text/event-stream; note=\"x,y;z\"")); + assertFalse(HttpUtils.isTextEventStreamContentType("text/event-stream; note=\"unterminated")); + } +} From b4c37a5159ae884e47a71baf2be1aedb4772eb73 Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Wed, 19 Aug 2026 11:23:57 +0800 Subject: [PATCH 2/8] Close streaming responses on body termination Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daa3657a-91d5-48fa-b7ae-6983091cab6a --- .../http/rest/AsyncRestProxy.java | 42 +++++++++++-------- .../http/rest/RestProxyBase.java | 35 ++++++++++++++-- .../http/rest/AsyncRestProxyTests.java | 38 +++++++++++++++-- 3 files changed, 89 insertions(+), 26 deletions(-) diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java index 1bf161a9431f..0211c9e92746 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java @@ -147,28 +147,28 @@ private Mono handleRestResponseReturnType(final HttpResponseDecoder.HttpDecod final boolean shouldPreserveResponseBodyAsStream = preserveResponseBodyAsStream || HttpUtils.isTextEventStreamContentType( response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); + final ResponseBodyOwner responseBodyOwner + = shouldPreserveResponseBodyAsStream ? new ResponseBodyOwner(response.getSourceResponse()) : null; if (methodParser.isStreamResponse()) { return Mono.fromSupplier(() -> new StreamResponse(response.getSourceResponse())); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) { final Type bodyType = TypeUtil.getRestResponseBodyType(entityType); if (TypeUtil.isTypeOrSubTypeOf(bodyType, Void.class)) { - return response.getSourceResponse() - .getBody() - .ignoreElements() - .then(Mono.fromCallable( - () -> createResponse(response, entityType, null, shouldPreserveResponseBodyAsStream))); + Flux responseBody + = responseBodyOwner == null ? response.getSourceResponse().getBody() : responseBodyOwner.getBody(); + return responseBody.ignoreElements() + .then(Mono.fromCallable(() -> createResponse(response, entityType, null, responseBodyOwner))); } else { return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType, - shouldPreserveResponseBodyAsStream) - .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject, - shouldPreserveResponseBodyAsStream)) - .switchIfEmpty(Mono.fromCallable( - () -> createResponse(response, entityType, null, shouldPreserveResponseBodyAsStream))); + responseBodyOwner) + .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject, responseBodyOwner)) + .switchIfEmpty( + Mono.fromCallable(() -> createResponse(response, entityType, null, responseBodyOwner))); } } else { // For now, we're just throwing if the Maybe didn't emit a value. return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType, - shouldPreserveResponseBodyAsStream); + responseBodyOwner); } } @@ -186,10 +186,17 @@ private static Function> decodeBytes(HttpResponseDecoder.Ht } static Mono handleBodyReturnType(HttpResponse sourceResponse, Function> getDecodedBody, - SwaggerMethodParser methodParser, Type entityType, boolean responseBodyStreaming) { + SwaggerMethodParser methodParser, Type entityType, ResponseBodyOwner responseBodyOwner) { final int responseStatusCode = sourceResponse.getStatusCode(); final HttpMethod httpMethod = methodParser.getHttpMethod(); final Type returnValueWireType = methodParser.getReturnValueWireType(); + if (responseBodyOwner == null + && HttpUtils + .isTextEventStreamContentType(sourceResponse.getHeaders().getValue(HttpHeaderName.CONTENT_TYPE))) { + responseBodyOwner = new ResponseBodyOwner(sourceResponse); + } + final Flux responseBody + = responseBodyOwner == null ? sourceResponse.getBody() : responseBodyOwner.getBody(); final Mono asyncResult; if (httpMethod == HttpMethod.HEAD @@ -208,20 +215,19 @@ static Mono handleBodyReturnType(HttpResponse sourceResponse, Function> - asyncResult = Mono.just(sourceResponse.getBody()); + asyncResult = Mono.just(responseBody); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, BinaryData.class)) { - String contentType = sourceResponse.getHeaders().getValue(HttpHeaderName.CONTENT_TYPE); // Mono // The raw response is directly used to create an instance of BinaryData which then provides // different methods to read the response. The reading of the response is delayed until BinaryData // is read and depending on which format the content is converted into, the response is not necessarily // fully copied into memory resulting in lesser overall memory usage. - if (responseBodyStreaming || HttpUtils.isTextEventStreamContentType(contentType)) { + if (responseBodyOwner != null) { // If the request or response content type identifies a stream, create a BinaryData instance with // bufferContent set to false. - asyncResult = BinaryData.fromFlux(sourceResponse.getBody(), null, false); + asyncResult = BinaryData.fromFlux(responseBody, null, false); } else { - asyncResult = BinaryData.fromFlux(sourceResponse.getBody()); + asyncResult = BinaryData.fromFlux(responseBody); } } else if (TypeUtil.isTypeOrSubTypeOf(entityType, InputStream.class)) { // Corresponds to the Open API 2.0 type "file" which is mapped to an InputStream. @@ -235,7 +241,7 @@ static Mono handleBodyReturnType(HttpResponse sourceResponse, Function handleBodyReturnType(HttpResponse sourceResponse, Function> getDecodedBody, SwaggerMethodParser methodParser, Type entityType) { - return handleBodyReturnType(sourceResponse, getDecodedBody, methodParser, entityType, false); + return handleBodyReturnType(sourceResponse, getDecodedBody, methodParser, entityType, null); } /** diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java index 2c70a09b1106..523c8100df6c 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java @@ -36,12 +36,14 @@ import com.azure.core.util.serializer.SerializerAdapter; import com.azure.core.util.tracing.Tracer; import reactor.core.Exceptions; +import reactor.core.publisher.Flux; import java.io.Closeable; import java.io.IOException; import java.lang.reflect.Method; import java.lang.reflect.Type; import java.net.URL; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.EnumSet; import java.util.concurrent.atomic.AtomicBoolean; @@ -184,6 +186,13 @@ public abstract void updateRequest(RequestDataConfiguration requestDataConfigura @SuppressWarnings({ "unchecked", "rawtypes" }) public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, Type entityType, Object bodyAsObject, boolean responseBodyStreaming) { + return createResponse(response, entityType, bodyAsObject, + responseBodyStreaming ? new ResponseBodyOwner(response.getSourceResponse()) : null); + } + + @SuppressWarnings({ "unchecked", "rawtypes" }) + final Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, Type entityType, + Object bodyAsObject, ResponseBodyOwner responseBodyOwner) { final Class> cls = (Class>) TypeUtil.getRawClass(entityType); final HttpResponse httpResponse = response.getSourceResponse(); @@ -197,8 +206,8 @@ public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, // If the type is either the Response or PagedResponse interface from azure-core a new instance of either // ResponseBase or PagedResponseBase can be returned. if (cls.equals(Response.class)) { - if (responseBodyStreaming) { - return cls.cast(new CloseableResponse<>(httpResponse, bodyAsObject)); + if (responseBodyOwner != null) { + return cls.cast(new CloseableResponse<>(responseBodyOwner, bodyAsObject)); } // For Response return a new instance of ResponseBase cast to the class. @@ -230,14 +239,32 @@ public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, } private static final class CloseableResponse extends SimpleResponse implements Closeable { + private final ResponseBodyOwner responseBodyOwner; + + private CloseableResponse(ResponseBodyOwner responseBodyOwner, T value) { + super(responseBodyOwner.response.getRequest(), responseBodyOwner.response.getStatusCode(), + responseBodyOwner.response.getHeaders(), value); + this.responseBodyOwner = responseBodyOwner; + } + + @Override + public void close() { + responseBodyOwner.close(); + } + } + + static final class ResponseBodyOwner implements Closeable { private final AtomicBoolean closed = new AtomicBoolean(); private final HttpResponse response; - private CloseableResponse(HttpResponse response, T value) { - super(response.getRequest(), response.getStatusCode(), response.getHeaders(), value); + ResponseBodyOwner(HttpResponse response) { this.response = response; } + Flux getBody() { + return Flux.using(() -> this, ignored -> response.getBody(), ResponseBodyOwner::close); + } + @Override public void close() { if (closed.compareAndSet(false, true)) { diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java index de451efe2e1e..fc578e9b4481 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java @@ -35,7 +35,7 @@ import java.lang.reflect.Type; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import static com.azure.core.CoreTestUtils.assertArraysEqual; @@ -213,12 +213,12 @@ public static Stream getResponseHeaderAndReplayability() { @MethodSource("streamingResponseOwnershipSupplier") public void streamingResponseIsUnbufferedAndCloseable(String accept, String contentType) throws IOException { byte[] expectedBytes = "hello".getBytes(StandardCharsets.UTF_8); - AtomicBoolean responseClosed = new AtomicBoolean(); + AtomicInteger responseCloseCount = new AtomicInteger(); HttpClient client = request -> Mono.just(new MockHttpResponse(request, 200, new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, contentType), expectedBytes) { @Override public void close() { - responseClosed.set(true); + responseCloseCount.incrementAndGet(); super.close(); } }); @@ -233,8 +233,38 @@ public void close() { assertTrue(response instanceof Closeable); assertFalse(response.getValue().isReplayable()); assertArraysEqual(expectedBytes, response.getValue().toBytes()); + assertEquals(1, responseCloseCount.get()); ((Closeable) response).close(); - assertTrue(responseClosed.get()); + assertEquals(1, responseCloseCount.get()); + } + + @Test + public void cancellingStreamingBinaryDataClosesResponse() { + byte[] expectedBytes = "hello".getBytes(StandardCharsets.UTF_8); + AtomicInteger responseCloseCount = new AtomicInteger(); + HttpClient client = request -> Mono.just(new MockHttpResponse(request, 200, + new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, "text/event-stream"), expectedBytes) { + @Override + public Flux getBody() { + return Flux.concat(Flux.just(ByteBuffer.wrap(expectedBytes)), Flux.never()); + } + + @Override + public void close() { + responseCloseCount.incrementAndGet(); + super.close(); + } + }); + MockService service = RestProxy.create(MockService.class, new HttpPipelineBuilder().httpClient(client).build()); + + Flux body = service.getStreamingResponse(new RequestOptions(), Context.NONE) + .flatMapMany(response -> response.getValue().toFluxByteBuffer()); + + StepVerifier.create(body) + .assertNext(buffer -> assertArraysEqual(expectedBytes, buffer.array())) + .thenCancel() + .verify(); + assertEquals(1, responseCloseCount.get()); } private static Stream streamingResponseOwnershipSupplier() { From 3f7229d2b8d25d0477534cd369bda8dcb4a8f33f Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Wed, 19 Aug 2026 13:05:50 +0800 Subject: [PATCH 3/8] Close sync BinaryData responses on termination Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c6a36a59-cfcf-4a11-b7d3-967de8071bd3 --- .../http/rest/SyncRestProxy.java | 39 ++++++++- .../http/rest/SyncRestProxyTests.java | 86 +++++++++++++------ 2 files changed, 95 insertions(+), 30 deletions(-) diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index 148f34820634..fc9bf3ee9225 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -22,6 +22,7 @@ import com.azure.core.util.serializer.SerializerEncoding; import com.azure.json.JsonSerializable; import com.azure.xml.XmlSerializable; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.io.IOException; @@ -30,6 +31,7 @@ import java.lang.reflect.Type; import java.nio.ByteBuffer; import java.util.EnumSet; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import static com.azure.core.implementation.ReflectionSerializable.serializeJsonSerializableToBytes; @@ -152,8 +154,10 @@ private Object handleRestResponseReturnType(HttpResponseDecoder.HttpDecodedRespo response.getSourceResponse().close(); return createResponse(response, entityType, null, responseBodyStreaming); } else { - Object bodyAsObject = handleBodyReturnType(response, methodParser, bodyType); - Response httpResponse = createResponse(response, entityType, bodyAsObject, responseBodyStreaming); + Object bodyAsObject = TypeUtil.isTypeOrSubTypeOf(bodyType, BinaryData.class) + ? getOwnedResponseBody(response.getSourceResponse()) + : handleBodyReturnType(response, methodParser, bodyType); + Response httpResponse = createResponse(response, entityType, bodyAsObject); if (httpResponse == null) { return createResponse(response, entityType, null, responseBodyStreaming); } @@ -202,6 +206,37 @@ private Object handleBodyReturnType(HttpResponseDecoder.HttpDecodedResponse resp return result; } + private static BinaryData getOwnedResponseBody(HttpResponse response) { + BinaryData responseBody = response.getBodyAsBinaryData(); + if (responseBody == null || responseBody.isReplayable()) { + return responseBody; + } + + ResponseBodyOwner responseBodyOwner = new ResponseBodyOwner(response); + return BinaryData + .fromFlux(responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), false) + .block(); + } + + private static final class ResponseBodyOwner { + private final AtomicBoolean closed = new AtomicBoolean(); + private final HttpResponse response; + + private ResponseBodyOwner(HttpResponse response) { + this.response = response; + } + + private Flux getBody(Flux responseBody) { + return Flux.using(() -> this, ignored -> responseBody, ResponseBodyOwner::close); + } + + private void close() { + if (closed.compareAndSet(false, true)) { + response.close(); + } + } + } + /** * Handle the provided asynchronous HTTP response and return the deserialized value. * diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java index 99809e7b1e63..62b0e8aee6ae 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java @@ -13,7 +13,6 @@ import com.azure.core.annotation.ServiceInterface; import com.azure.core.http.HttpClient; import com.azure.core.http.HttpHeaderName; -import com.azure.core.http.HttpHeaders; import com.azure.core.http.HttpMethod; import com.azure.core.http.HttpPipeline; import com.azure.core.http.HttpPipelineBuilder; @@ -30,19 +29,24 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import reactor.core.Disposable; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; import java.io.ByteArrayInputStream; -import java.io.Closeable; import java.io.IOException; import java.io.InputStream; +import java.nio.ByteBuffer; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; /** @@ -72,7 +76,7 @@ Response testMethod(@BodyParam("application/octet-stream") BinaryData data @Get("my/url/path") @ExpectedResponses({ 200 }) - Response getStreamingResponse(RequestOptions options, Context context); + Response testBinaryDataResponse(Context context); } @Test @@ -199,10 +203,49 @@ public void testInputStream() throws IOException { assertEquals("hello", new String(bytes)); } - @ParameterizedTest - @MethodSource("streamingResponseOwnershipSupplier") - public void streamingResponseIsCloseable(String accept, String contentType) throws IOException { - AtomicBoolean responseClosed = new AtomicBoolean(); + @Test + public void binaryDataResponseClosesOnCompletion() { + AtomicInteger responseCloseCount = new AtomicInteger(); + TestInterface testInterface + = createBinaryDataService(Flux.just(ByteBuffer.wrap("hello".getBytes())), responseCloseCount); + + BinaryData responseBody = testInterface.testBinaryDataResponse(Context.NONE).getValue(); + + assertFalse(responseBody.isReplayable()); + assertEquals("hello", responseBody.toString()); + assertEquals(1, responseCloseCount.get()); + assertEquals("hello", responseBody.toString()); + assertEquals(1, responseCloseCount.get()); + } + + @Test + public void binaryDataResponseClosesOnCancellation() { + AtomicInteger responseCloseCount = new AtomicInteger(); + Flux responseBody = Flux.concat(Flux.just(ByteBuffer.wrap("hello".getBytes())), Flux.never()); + TestInterface testInterface = createBinaryDataService(responseBody, responseCloseCount); + + Disposable subscription + = testInterface.testBinaryDataResponse(Context.NONE).getValue().toFluxByteBuffer().subscribe(); + subscription.dispose(); + + assertEquals(1, responseCloseCount.get()); + } + + @Test + public void binaryDataResponseClosesOnError() { + AtomicInteger responseCloseCount = new AtomicInteger(); + TestInterface testInterface + = createBinaryDataService(Flux.error(new IllegalStateException("Body read failed.")), responseCloseCount); + + StepVerifier.create(testInterface.testBinaryDataResponse(Context.NONE).getValue().toFluxByteBuffer()) + .expectErrorMessage("Body read failed.") + .verify(); + + assertEquals(1, responseCloseCount.get()); + } + + private static TestInterface createBinaryDataService(Flux responseBody, + AtomicInteger responseCloseCount) { HttpClient client = new HttpClient() { @Override public Mono send(HttpRequest request) { @@ -211,34 +254,21 @@ public Mono send(HttpRequest request) { @Override public HttpResponse sendSync(HttpRequest request, Context context) { - return new MockHttpResponse(request, 200, - new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, contentType), "hello".getBytes()) { + return new MockHttpResponse(request, 200) { + @Override + public BinaryData getBodyAsBinaryData() { + return BinaryData.fromFlux(responseBody, null, false).block(); + } + @Override public void close() { - responseClosed.set(true); + responseCloseCount.incrementAndGet(); super.close(); } }; } }; - TestInterface service - = RestProxy.create(TestInterface.class, new HttpPipelineBuilder().httpClient(client).build()); - RequestOptions options = new RequestOptions(); - if (accept != null) { - options.setHeader(HttpHeaderName.ACCEPT, accept); - } - - Response response = service.getStreamingResponse(options, Context.NONE); - - assertTrue(response instanceof Closeable); - assertEquals("hello", response.getValue().toString()); - ((Closeable) response).close(); - assertTrue(responseClosed.get()); - } - - private static Stream streamingResponseOwnershipSupplier() { - return Stream.of(Arguments.of("text/event-stream", "application/json"), - Arguments.of(null, "text/event-stream; charset=utf-8")); + return RestProxy.create(TestInterface.class, new HttpPipelineBuilder().httpClient(client).build()); } private static Stream mergeRequestOptionsContextSupplier() { From 46924a390cfde76a7052f3f1b60593d973895a7f Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Wed, 19 Aug 2026 13:15:39 +0800 Subject: [PATCH 4/8] Keep streaming response ownership in body Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c6a36a59-cfcf-4a11-b7d3-967de8071bd3 --- .../http/rest/AsyncRestProxy.java | 8 ++--- .../http/rest/RestProxyBase.java | 36 ++++--------------- .../http/rest/SyncRestProxy.java | 35 +++--------------- .../http/rest/AsyncRestProxyTests.java | 6 +--- 4 files changed, 14 insertions(+), 71 deletions(-) diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java index 0211c9e92746..4f967f56741c 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java @@ -157,13 +157,11 @@ private Mono handleRestResponseReturnType(final HttpResponseDecoder.HttpDecod Flux responseBody = responseBodyOwner == null ? response.getSourceResponse().getBody() : responseBodyOwner.getBody(); return responseBody.ignoreElements() - .then(Mono.fromCallable(() -> createResponse(response, entityType, null, responseBodyOwner))); + .then(Mono.fromCallable(() -> createResponse(response, entityType, null))); } else { return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType, - responseBodyOwner) - .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject, responseBodyOwner)) - .switchIfEmpty( - Mono.fromCallable(() -> createResponse(response, entityType, null, responseBodyOwner))); + responseBodyOwner).map(bodyAsObject -> createResponse(response, entityType, bodyAsObject)) + .switchIfEmpty(Mono.fromCallable(() -> createResponse(response, entityType, null))); } } else { // For now, we're just throwing if the Maybe didn't emit a value. diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java index 523c8100df6c..cd9441bb5548 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java @@ -22,7 +22,6 @@ import com.azure.core.http.rest.RequestOptions; import com.azure.core.http.rest.Response; import com.azure.core.http.rest.ResponseBase; -import com.azure.core.http.rest.SimpleResponse; import com.azure.core.implementation.ReflectiveInvoker; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.http.UnexpectedExceptionInformation; @@ -178,21 +177,13 @@ public abstract void updateRequest(RequestDataConfiguration requestDataConfigura * @param response the decoded response * @param entityType the type of the response entity * @param bodyAsObject the response body as an object - * @param responseBodyStreaming whether the response owns an unconsumed streaming body * @return the {@link Response} * @throws RuntimeException If the response type is a PagedResponse and the bodyAsObject is not an instance of * Page. */ @SuppressWarnings({ "unchecked", "rawtypes" }) public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, Type entityType, - Object bodyAsObject, boolean responseBodyStreaming) { - return createResponse(response, entityType, bodyAsObject, - responseBodyStreaming ? new ResponseBodyOwner(response.getSourceResponse()) : null); - } - - @SuppressWarnings({ "unchecked", "rawtypes" }) - final Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, Type entityType, - Object bodyAsObject, ResponseBodyOwner responseBodyOwner) { + Object bodyAsObject) { final Class> cls = (Class>) TypeUtil.getRawClass(entityType); final HttpResponse httpResponse = response.getSourceResponse(); @@ -206,10 +197,6 @@ final Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, // If the type is either the Response or PagedResponse interface from azure-core a new instance of either // ResponseBase or PagedResponseBase can be returned. if (cls.equals(Response.class)) { - if (responseBodyOwner != null) { - return cls.cast(new CloseableResponse<>(responseBodyOwner, bodyAsObject)); - } - // For Response return a new instance of ResponseBase cast to the class. return cls.cast(new ResponseBase<>(request, statusCode, headers, bodyAsObject, decodedHeaders)); } else if (cls.equals(PagedResponse.class)) { @@ -238,21 +225,6 @@ final Response createResponse(HttpResponseDecoder.HttpDecodedResponse response, return RESPONSE_CONSTRUCTORS_CACHE.invoke(constructorReflectiveInvoker, response, bodyAsObject); } - private static final class CloseableResponse extends SimpleResponse implements Closeable { - private final ResponseBodyOwner responseBodyOwner; - - private CloseableResponse(ResponseBodyOwner responseBodyOwner, T value) { - super(responseBodyOwner.response.getRequest(), responseBodyOwner.response.getStatusCode(), - responseBodyOwner.response.getHeaders(), value); - this.responseBodyOwner = responseBodyOwner; - } - - @Override - public void close() { - responseBodyOwner.close(); - } - } - static final class ResponseBodyOwner implements Closeable { private final AtomicBoolean closed = new AtomicBoolean(); private final HttpResponse response; @@ -262,7 +234,11 @@ static final class ResponseBodyOwner implements Closeable { } Flux getBody() { - return Flux.using(() -> this, ignored -> response.getBody(), ResponseBodyOwner::close); + return getBody(response.getBody()); + } + + Flux getBody(Flux responseBody) { + return Flux.using(() -> this, ignored -> responseBody, ResponseBodyOwner::close); } @Override diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index fc9bf3ee9225..a6a54fcef863 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -3,7 +3,6 @@ package com.azure.core.implementation.http.rest; -import com.azure.core.http.HttpHeaderName; import com.azure.core.http.HttpMethod; import com.azure.core.http.HttpPipeline; import com.azure.core.http.HttpRequest; @@ -14,7 +13,6 @@ import com.azure.core.implementation.ImplUtils; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.serializer.HttpResponseDecoder; -import com.azure.core.implementation.util.HttpUtils; import com.azure.core.util.Base64Url; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; @@ -22,7 +20,6 @@ import com.azure.core.util.serializer.SerializerEncoding; import com.azure.json.JsonSerializable; import com.azure.xml.XmlSerializable; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.io.IOException; @@ -31,7 +28,6 @@ import java.lang.reflect.Type; import java.nio.ByteBuffer; import java.util.EnumSet; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import static com.azure.core.implementation.ReflectionSerializable.serializeJsonSerializableToBytes; @@ -142,24 +138,21 @@ private HttpResponseDecoder.HttpDecodedResponse ensureExpectedStatus( } private Object handleRestResponseReturnType(HttpResponseDecoder.HttpDecodedResponse response, - SwaggerMethodParser methodParser, Type entityType, boolean responseBodyStreaming) { - responseBodyStreaming = responseBodyStreaming - || HttpUtils.isTextEventStreamContentType( - response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); + SwaggerMethodParser methodParser, Type entityType) { if (methodParser.isStreamResponse()) { return new StreamResponse(response.getSourceResponse()); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) { final Type bodyType = TypeUtil.getRestResponseBodyType(entityType); if (TypeUtil.isTypeOrSubTypeOf(bodyType, Void.class)) { response.getSourceResponse().close(); - return createResponse(response, entityType, null, responseBodyStreaming); + return createResponse(response, entityType, null); } else { Object bodyAsObject = TypeUtil.isTypeOrSubTypeOf(bodyType, BinaryData.class) ? getOwnedResponseBody(response.getSourceResponse()) : handleBodyReturnType(response, methodParser, bodyType); Response httpResponse = createResponse(response, entityType, bodyAsObject); if (httpResponse == null) { - return createResponse(response, entityType, null, responseBodyStreaming); + return createResponse(response, entityType, null); } return httpResponse; } @@ -218,25 +211,6 @@ private static BinaryData getOwnedResponseBody(HttpResponse response) { .block(); } - private static final class ResponseBodyOwner { - private final AtomicBoolean closed = new AtomicBoolean(); - private final HttpResponse response; - - private ResponseBodyOwner(HttpResponse response) { - this.response = response; - } - - private Flux getBody(Flux responseBody) { - return Flux.using(() -> this, ignored -> responseBody, ResponseBodyOwner::close); - } - - private void close() { - if (closed.compareAndSet(false, true)) { - response.close(); - } - } - } - /** * Handle the provided asynchronous HTTP response and return the deserialized value. * @@ -260,8 +234,7 @@ private Object handleRestReturnType(HttpResponseDecoder.HttpDecodedResponse http } else { // ProxyMethod ReturnType: T where T != async (Mono, Flux) or sync Void // Block the deserialization until a value T is received - result = handleRestResponseReturnType(httpDecodedResponse, methodParser, returnType, - HttpUtils.shouldPreserveResponseBodyAsStream(context)); + result = handleRestResponseReturnType(httpDecodedResponse, methodParser, returnType); } return result; } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java index fc578e9b4481..1b5244a72e6c 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java @@ -29,7 +29,6 @@ import reactor.core.publisher.Mono; import reactor.test.StepVerifier; -import java.io.Closeable; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Type; @@ -211,7 +210,7 @@ public static Stream getResponseHeaderAndReplayability() { @ParameterizedTest @MethodSource("streamingResponseOwnershipSupplier") - public void streamingResponseIsUnbufferedAndCloseable(String accept, String contentType) throws IOException { + public void streamingResponseIsUnbufferedAndClosedOnConsumption(String accept, String contentType) { byte[] expectedBytes = "hello".getBytes(StandardCharsets.UTF_8); AtomicInteger responseCloseCount = new AtomicInteger(); HttpClient client = request -> Mono.just(new MockHttpResponse(request, 200, @@ -230,12 +229,9 @@ public void close() { Response response = service.getStreamingResponse(options, Context.NONE).block(); - assertTrue(response instanceof Closeable); assertFalse(response.getValue().isReplayable()); assertArraysEqual(expectedBytes, response.getValue().toBytes()); assertEquals(1, responseCloseCount.get()); - ((Closeable) response).close(); - assertEquals(1, responseCloseCount.get()); } @Test From ebcf12c7fee9a6cb4dd5e6a86317ef0b7bc9fa52 Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Thu, 20 Aug 2026 14:43:28 +0800 Subject: [PATCH 5/8] Preserve synchronous BinaryData streaming Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 607c904a-42fc-440a-b53d-dbb61a5cbd4e --- .../http/rest/StreamingResponseContent.java | 119 ++++++++++++++++++ .../http/rest/SyncRestProxy.java | 6 +- .../http/rest/SyncRestProxyTests.java | 57 +++++++++ 3 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java new file mode 100644 index 000000000000..bdb40e0cc668 --- /dev/null +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.core.implementation.http.rest; + +import com.azure.core.implementation.FluxInputStream; +import com.azure.core.implementation.util.BinaryDataContent; +import com.azure.core.implementation.util.BinaryDataContentType; +import com.azure.core.implementation.util.FluxByteBufferContent; +import com.azure.core.util.serializer.ObjectSerializer; +import com.azure.core.util.serializer.TypeReference; +import com.azure.json.JsonWriter; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.AsynchronousByteChannel; +import java.nio.channels.WritableByteChannel; + +/** + * Non-replayable response content that exposes both reactive and synchronous streaming access. + */ +final class StreamingResponseContent extends BinaryDataContent { + private final FluxByteBufferContent content; + private final RestProxyBase.ResponseBodyOwner responseBodyOwner; + + StreamingResponseContent(Flux content, Long length, RestProxyBase.ResponseBodyOwner responseBodyOwner) { + this.content = new FluxByteBufferContent(content, length, false); + this.responseBodyOwner = responseBodyOwner; + } + + @Override + public Long getLength() { + return content.getLength(); + } + + @Override + public String toString() { + return content.toString(); + } + + @Override + public byte[] toBytes() { + return content.toBytes(); + } + + @Override + public T toObject(TypeReference typeReference, ObjectSerializer serializer) { + return content.toObject(typeReference, serializer); + } + + @Override + public InputStream toStream() { + return new FilterInputStream(new FluxInputStream(content.toFluxByteBuffer())) { + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + responseBodyOwner.close(); + } + } + }; + } + + @Override + public ByteBuffer toByteBuffer() { + return content.toByteBuffer(); + } + + @Override + public Flux toFluxByteBuffer() { + return content.toFluxByteBuffer(); + } + + @Override + public void writeTo(OutputStream outputStream) throws IOException { + content.writeTo(outputStream); + } + + @Override + public void writeTo(WritableByteChannel channel) throws IOException { + content.writeTo(channel); + } + + @Override + public Mono writeTo(AsynchronousByteChannel channel) { + return content.writeTo(channel); + } + + @Override + public void writeTo(JsonWriter jsonWriter) throws IOException { + content.writeTo(jsonWriter); + } + + @Override + public boolean isReplayable() { + return false; + } + + @Override + public BinaryDataContent toReplayableContent() { + return content.toReplayableContent(); + } + + @Override + public Mono toReplayableContentAsync() { + return content.toReplayableContentAsync(); + } + + @Override + public BinaryDataContentType getContentType() { + return content.getContentType(); + } +} diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index a6a54fcef863..f09c88be2d75 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -13,6 +13,7 @@ import com.azure.core.implementation.ImplUtils; import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.serializer.HttpResponseDecoder; +import com.azure.core.implementation.util.BinaryDataHelper; import com.azure.core.util.Base64Url; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; @@ -206,9 +207,8 @@ private static BinaryData getOwnedResponseBody(HttpResponse response) { } ResponseBodyOwner responseBodyOwner = new ResponseBodyOwner(response); - return BinaryData - .fromFlux(responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), false) - .block(); + return BinaryDataHelper.createBinaryData(new StreamingResponseContent( + responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), responseBodyOwner)); } /** diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java index 62b0e8aee6ae..090abf773398 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java @@ -38,6 +38,7 @@ import java.io.IOException; import java.io.InputStream; import java.nio.ByteBuffer; +import java.time.Duration; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -47,6 +48,9 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; /** @@ -231,6 +235,59 @@ public void binaryDataResponseClosesOnCancellation() { assertEquals(1, responseCloseCount.get()); } + @Test + public void binaryDataResponseToStreamReadsWithoutWaitingForCompletion() throws IOException { + byte[] firstChunk = new byte[8192]; + byte[] secondChunk = new byte[8192]; + firstChunk[0] = 1; + secondChunk[0] = 2; + AtomicInteger responseCloseCount = new AtomicInteger(); + Flux responseBody + = Flux.concat(Flux.just(ByteBuffer.wrap(firstChunk), ByteBuffer.wrap(secondChunk)), Flux.never()); + TestInterface testInterface = createBinaryDataService(responseBody, responseCloseCount); + + BinaryData binaryData = testInterface.testBinaryDataResponse(Context.NONE).getValue(); + assertFalse(binaryData.isReplayable()); + + assertTimeoutPreemptively(Duration.ofSeconds(5), () -> { + try (InputStream stream = binaryData.toStream()) { + byte[] actualFirstChunk = new byte[firstChunk.length]; + byte[] actualSecondChunk = new byte[secondChunk.length]; + assertEquals(firstChunk.length, stream.read(actualFirstChunk)); + assertEquals(secondChunk.length, stream.read(actualSecondChunk)); + assertArrayEquals(firstChunk, actualFirstChunk); + assertArrayEquals(secondChunk, actualSecondChunk); + assertEquals(0, responseCloseCount.get()); + } + }); + + assertEquals(1, responseCloseCount.get()); + } + + @Test + public void binaryDataResponseToStreamClosesOnError() throws IOException { + AtomicInteger responseCloseCount = new AtomicInteger(); + TestInterface testInterface + = createBinaryDataService(Flux.error(new IllegalStateException("Body read failed.")), responseCloseCount); + + try (InputStream stream = testInterface.testBinaryDataResponse(Context.NONE).getValue().toStream()) { + assertThrows(IOException.class, stream::read); + assertEquals(1, responseCloseCount.get()); + } + + assertEquals(1, responseCloseCount.get()); + } + + @Test + public void binaryDataResponseToStreamClosesBeforeFirstRead() throws IOException { + AtomicInteger responseCloseCount = new AtomicInteger(); + TestInterface testInterface = createBinaryDataService(Flux.never(), responseCloseCount); + + testInterface.testBinaryDataResponse(Context.NONE).getValue().toStream().close(); + + assertEquals(1, responseCloseCount.get()); + } + @Test public void binaryDataResponseClosesOnError() { AtomicInteger responseCloseCount = new AtomicInteger(); From 2077a38947ecde55b48a49824a0f7d95c461edb4 Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Thu, 20 Aug 2026 19:24:45 +0800 Subject: [PATCH 6/8] Stream non-replayable Flux BinaryData Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 607c904a-42fc-440a-b53d-dbb61a5cbd4e --- .../core/implementation/FluxInputStream.java | 11 +- .../http/rest/StreamingResponseContent.java | 119 ------------------ .../http/rest/SyncRestProxy.java | 5 +- .../util/FluxByteBufferContent.java | 12 +- .../implementation/FluxInputStreamTests.java | 100 +++++++++++++++ .../http/rest/SyncRestProxyTests.java | 7 +- .../com/azure/core/util/BinaryDataTest.java | 19 +++ 7 files changed, 146 insertions(+), 127 deletions(-) delete mode 100644 sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java index aa962955a1b2..a5d3a80269b0 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java @@ -29,7 +29,7 @@ public class FluxInputStream extends InputStream { private final Flux data; // Subscription to request more data from as needed - private Subscription subscription; + private volatile Subscription subscription; private ByteArrayInputStream buffer; @@ -151,6 +151,9 @@ public void close() throws IOException { // Unblock any thread waiting in blockForData(). lock.lock(); try { + if (!subscribed) { + subscribeToData(); + } waitingForData = false; dataAvailable.signal(); } finally { @@ -227,13 +230,13 @@ private void subscribeToData() { this::signalOnCompleteOrError, // Subscription consumer subscription -> { + this.subscription = subscription; + this.subscribed = true; if (this.closed) { subscription.cancel(); return; } - this.subscription = subscription; - this.subscribed = true; - this.subscription.request(1); + subscription.request(1); }); } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java deleted file mode 100644 index bdb40e0cc668..000000000000 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/StreamingResponseContent.java +++ /dev/null @@ -1,119 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.core.implementation.http.rest; - -import com.azure.core.implementation.FluxInputStream; -import com.azure.core.implementation.util.BinaryDataContent; -import com.azure.core.implementation.util.BinaryDataContentType; -import com.azure.core.implementation.util.FluxByteBufferContent; -import com.azure.core.util.serializer.ObjectSerializer; -import com.azure.core.util.serializer.TypeReference; -import com.azure.json.JsonWriter; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import java.io.FilterInputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; -import java.nio.ByteBuffer; -import java.nio.channels.AsynchronousByteChannel; -import java.nio.channels.WritableByteChannel; - -/** - * Non-replayable response content that exposes both reactive and synchronous streaming access. - */ -final class StreamingResponseContent extends BinaryDataContent { - private final FluxByteBufferContent content; - private final RestProxyBase.ResponseBodyOwner responseBodyOwner; - - StreamingResponseContent(Flux content, Long length, RestProxyBase.ResponseBodyOwner responseBodyOwner) { - this.content = new FluxByteBufferContent(content, length, false); - this.responseBodyOwner = responseBodyOwner; - } - - @Override - public Long getLength() { - return content.getLength(); - } - - @Override - public String toString() { - return content.toString(); - } - - @Override - public byte[] toBytes() { - return content.toBytes(); - } - - @Override - public T toObject(TypeReference typeReference, ObjectSerializer serializer) { - return content.toObject(typeReference, serializer); - } - - @Override - public InputStream toStream() { - return new FilterInputStream(new FluxInputStream(content.toFluxByteBuffer())) { - @Override - public void close() throws IOException { - try { - super.close(); - } finally { - responseBodyOwner.close(); - } - } - }; - } - - @Override - public ByteBuffer toByteBuffer() { - return content.toByteBuffer(); - } - - @Override - public Flux toFluxByteBuffer() { - return content.toFluxByteBuffer(); - } - - @Override - public void writeTo(OutputStream outputStream) throws IOException { - content.writeTo(outputStream); - } - - @Override - public void writeTo(WritableByteChannel channel) throws IOException { - content.writeTo(channel); - } - - @Override - public Mono writeTo(AsynchronousByteChannel channel) { - return content.writeTo(channel); - } - - @Override - public void writeTo(JsonWriter jsonWriter) throws IOException { - content.writeTo(jsonWriter); - } - - @Override - public boolean isReplayable() { - return false; - } - - @Override - public BinaryDataContent toReplayableContent() { - return content.toReplayableContent(); - } - - @Override - public Mono toReplayableContentAsync() { - return content.toReplayableContentAsync(); - } - - @Override - public BinaryDataContentType getContentType() { - return content.getContentType(); - } -} diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index f09c88be2d75..f4269a0e0a4b 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -14,6 +14,7 @@ import com.azure.core.implementation.TypeUtil; import com.azure.core.implementation.serializer.HttpResponseDecoder; import com.azure.core.implementation.util.BinaryDataHelper; +import com.azure.core.implementation.util.FluxByteBufferContent; import com.azure.core.util.Base64Url; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; @@ -207,8 +208,8 @@ private static BinaryData getOwnedResponseBody(HttpResponse response) { } ResponseBodyOwner responseBodyOwner = new ResponseBodyOwner(response); - return BinaryDataHelper.createBinaryData(new StreamingResponseContent( - responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), responseBodyOwner)); + return BinaryDataHelper.createBinaryData(new FluxByteBufferContent( + responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), false)); } /** diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java index 1ccc4414be8e..cad28d1a76c9 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java @@ -3,6 +3,7 @@ package com.azure.core.implementation.util; +import com.azure.core.implementation.FluxInputStream; import com.azure.core.util.FluxUtil; import com.azure.core.util.logging.ClientLogger; import com.azure.core.util.serializer.ObjectSerializer; @@ -102,7 +103,16 @@ public T toObject(TypeReference typeReference, ObjectSerializer serialize @Override public InputStream toStream() { - return new ByteArrayInputStream(toBytes()); + byte[] cachedBytes = BYTES_UPDATER.get(this); + if (cachedBytes != null) { + return new ByteArrayInputStream(cachedBytes); + } + + if (isReplayable) { + return new ByteArrayInputStream(toBytes()); + } + + return new FluxInputStream(content); } @Override diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/FluxInputStreamTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/FluxInputStreamTests.java index 9408b4b83217..7ef7330252bd 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/FluxInputStreamTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/FluxInputStreamTests.java @@ -9,6 +9,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import org.reactivestreams.Subscription; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -19,10 +20,18 @@ import java.nio.charset.Charset; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; public class FluxInputStreamTests { private static final int KB = 1024; @@ -93,6 +102,85 @@ public void fluxInputStreamWithEmptyByteBuffers() throws IOException { } } + @Test + public void closeBeforeFirstReadSubscribesAndCancelsWithoutDemand() throws IOException { + AtomicInteger subscribeCalls = new AtomicInteger(); + AtomicInteger cancelCalls = new AtomicInteger(); + AtomicInteger disposeCalls = new AtomicInteger(); + AtomicLong requested = new AtomicLong(); + Flux data = Flux.using(Object::new, + ignored -> Flux.never() + .doOnSubscribe(subscription -> subscribeCalls.incrementAndGet()) + .doOnRequest(requested::addAndGet) + .doOnCancel(cancelCalls::incrementAndGet), + ignored -> disposeCalls.incrementAndGet()); + + FluxInputStream stream = new FluxInputStream(data); + stream.close(); + stream.close(); + + assertEquals(1, subscribeCalls.get()); + assertEquals(0, requested.get()); + assertEquals(1, cancelCalls.get()); + assertEquals(1, disposeCalls.get()); + } + + @Test + public void closeDuringFirstSubscriptionCancelsPublishedSubscription() throws Exception { + CountDownLatch subscribeEntered = new CountDownLatch(1); + AtomicInteger cancelCalls = new AtomicInteger(); + AtomicInteger disposeCalls = new AtomicInteger(); + AtomicLong requested = new AtomicLong(); + AtomicReference closeThreadReference = new AtomicReference<>(); + Flux data = Flux.using(Object::new, ignored -> Flux.from(subscriber -> { + subscribeEntered.countDown(); + awaitThreadWaiting(closeThreadReference); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long count) { + requested.addAndGet(count); + } + + @Override + public void cancel() { + cancelCalls.incrementAndGet(); + } + }); + }), ignored -> disposeCalls.incrementAndGet()); + FluxInputStream stream = new FluxInputStream(data); + AtomicReference readError = new AtomicReference<>(); + AtomicReference closeError = new AtomicReference<>(); + Thread readThread = new Thread(() -> { + try { + stream.read(); + } catch (Throwable throwable) { + readError.set(throwable); + } + }); + Thread closeThread = new Thread(() -> { + try { + stream.close(); + } catch (Throwable throwable) { + closeError.set(throwable); + } + }); + closeThreadReference.set(closeThread); + + readThread.start(); + assertTrue(subscribeEntered.await(5, TimeUnit.SECONDS)); + closeThread.start(); + readThread.join(TimeUnit.SECONDS.toMillis(5)); + closeThread.join(TimeUnit.SECONDS.toMillis(5)); + + assertFalse(readThread.isAlive()); + assertFalse(closeThread.isAlive()); + assertEquals(1, cancelCalls.get()); + assertEquals(0, requested.get()); + assertEquals(1, disposeCalls.get()); + assertNull(closeError.get()); + assertTrue(readError.get() instanceof IllegalStateException); + } + @ParameterizedTest @MethodSource("fluxInputStreamErrorSupplier") public void fluxInputStreamError(RuntimeException exception) { @@ -145,4 +233,16 @@ public Mono getBodyAsString(Charset charset) { new HttpResponseException("Mock exception", httpResponse, null), new UncheckedIOException(new IOException("Mock IO Exception."))); } + + private static void awaitThreadWaiting(AtomicReference threadReference) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (System.nanoTime() < deadline) { + Thread thread = threadReference.get(); + if (thread != null && thread.getState() == Thread.State.WAITING) { + return; + } + Thread.yield(); + } + throw new AssertionError("Close thread didn't wait for the stream lock."); + } } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java index 090abf773398..98099311266d 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java @@ -23,6 +23,8 @@ import com.azure.core.http.rest.Response; import com.azure.core.http.rest.RestProxy; import com.azure.core.http.rest.StreamResponse; +import com.azure.core.implementation.util.BinaryDataHelper; +import com.azure.core.implementation.util.FluxByteBufferContent; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import org.junit.jupiter.api.Test; @@ -49,6 +51,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -283,7 +286,9 @@ public void binaryDataResponseToStreamClosesBeforeFirstRead() throws IOException AtomicInteger responseCloseCount = new AtomicInteger(); TestInterface testInterface = createBinaryDataService(Flux.never(), responseCloseCount); - testInterface.testBinaryDataResponse(Context.NONE).getValue().toStream().close(); + BinaryData responseBody = testInterface.testBinaryDataResponse(Context.NONE).getValue(); + assertInstanceOf(FluxByteBufferContent.class, BinaryDataHelper.getContent(responseBody)); + responseBody.toStream().close(); assertEquals(1, responseCloseCount.get()); } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/util/BinaryDataTest.java b/sdk/core/azure-core/src/test/java/com/azure/core/util/BinaryDataTest.java index 187a09f9cd2d..3b03704e0aaf 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/util/BinaryDataTest.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/util/BinaryDataTest.java @@ -82,6 +82,7 @@ import static org.junit.jupiter.api.Assertions.assertNotSame; 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; /** @@ -567,6 +568,24 @@ public void fluxContent() { .verifyComplete(); } + @Test + public void nonReplayableFluxToStreamReadsWithoutWaitingForCompletion() { + byte[] expected = "hello".getBytes(StandardCharsets.UTF_8); + AtomicBoolean cancelled = new AtomicBoolean(); + Flux content + = Flux.concat(Flux.just(ByteBuffer.wrap(expected)), Flux.never()).doOnCancel(() -> cancelled.set(true)); + BinaryData data = BinaryData.fromFlux(content, null, false).block(); + + assertTimeoutPreemptively(Duration.ofSeconds(5), () -> { + try (InputStream stream = data.toStream()) { + byte[] actual = new byte[expected.length]; + assertEquals(expected.length, stream.read(actual)); + assertTrue(Arrays.equals(expected, actual)); + } + }); + assertTrue(cancelled.get()); + } + @Test public void testFromFile() throws Exception { Path file = Files.createTempFile("binaryDataFromFile" + UUID.randomUUID(), ".txt"); From 517f83681da6b9def22f8c2abe5041c20ac1b97d Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Thu, 20 Aug 2026 20:15:19 +0800 Subject: [PATCH 7/8] Document non-replayable Flux streaming Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 607c904a-42fc-440a-b53d-dbb61a5cbd4e --- sdk/core/azure-core/CHANGELOG.md | 2 ++ .../java/com/azure/core/implementation/FluxInputStream.java | 6 ++++++ .../core/implementation/util/FluxByteBufferContent.java | 6 ++++++ 3 files changed, 14 insertions(+) diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md index b6533eee3f70..a3faaedbe612 100644 --- a/sdk/core/azure-core/CHANGELOG.md +++ b/sdk/core/azure-core/CHANGELOG.md @@ -8,6 +8,8 @@ ### Bugs Fixed +- Fixed synchronous streaming of non-replayable `BinaryData` response bodies. + ### Other Changes ## 1.59.0 (2026-08-12) diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java index a5d3a80269b0..a5b06be1f9f6 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java @@ -141,6 +141,12 @@ public int read(byte[] b, int off, int len) throws IOException { } } + /** + * Closes the stream and cancels its Flux subscription. If the stream has not been read, closing subscribes and + * immediately cancels without requesting data so publisher cleanup can run. + * + * @throws IOException if the stream cannot be closed. + */ @Override public void close() throws IOException { closed = true; diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java index cad28d1a76c9..ca0ae3639a7e 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java @@ -101,6 +101,12 @@ public T toObject(TypeReference typeReference, ObjectSerializer serialize return serializer.deserializeFromBytes(toBytes(), typeReference); } + /** + * Returns an in-memory stream for cached or replayable content. Uncached non-replayable content is streamed + * incrementally without buffering the full Flux. + * + * @return A stream over this content. + */ @Override public InputStream toStream() { byte[] cachedBytes = BYTES_UPDATER.get(this); From ea1b57e9b3ccd79b782b830e581d29b7d3c57167 Mon Sep 17 00:00:00 2001 From: "Xiaofei Cao (from Dev Box)" Date: Fri, 21 Aug 2026 15:04:19 +0800 Subject: [PATCH 8/8] Use response content type for SSE Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 607c904a-42fc-440a-b53d-dbb61a5cbd4e --- .../http/rest/AsyncRestProxy.java | 26 ++++++------- .../http/rest/RestProxyBase.java | 8 ---- .../http/rest/SyncRestProxy.java | 2 - .../core/implementation/util/HttpUtils.java | 37 ------------------- .../http/policy/HttpLoggingPolicyTests.java | 12 ++---- .../http/rest/AsyncRestProxyTests.java | 12 +++--- .../implementation/util/HttpUtilsTests.java | 8 ---- 7 files changed, 21 insertions(+), 84 deletions(-) diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java index 4f967f56741c..fd89f5d98fc3 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java @@ -82,8 +82,6 @@ public Object invoke(Object proxy, Method method, RequestOptions options, EnumSe requestCallback.accept(request); } - context = updateRequestContext(request, context); - final Context finalContext = context; final Mono asyncResponse = RestProxyUtils.validateLengthAsync(request).flatMap(r -> { // correlates logs @@ -143,12 +141,11 @@ private Mono ensureExpectedStatus( } private Mono handleRestResponseReturnType(final HttpResponseDecoder.HttpDecodedResponse response, - final SwaggerMethodParser methodParser, final Type entityType, boolean preserveResponseBodyAsStream) { - final boolean shouldPreserveResponseBodyAsStream = preserveResponseBodyAsStream - || HttpUtils.isTextEventStreamContentType( - response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); + final SwaggerMethodParser methodParser, final Type entityType) { + final boolean isTextEventStream = HttpUtils.isTextEventStreamContentType( + response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE)); final ResponseBodyOwner responseBodyOwner - = shouldPreserveResponseBodyAsStream ? new ResponseBodyOwner(response.getSourceResponse()) : null; + = isTextEventStream ? new ResponseBodyOwner(response.getSourceResponse()) : null; if (methodParser.isStreamResponse()) { return Mono.fromSupplier(() -> new StreamResponse(response.getSourceResponse())); } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) { @@ -221,8 +218,8 @@ static Mono handleBodyReturnType(HttpResponse sourceResponse, Function errorOptionsSet) { final Mono asyncExpectedResponse = endSpanWhenDone( ensureExpectedStatus(asyncHttpDecodedResponse, methodParser, options, errorOptionsSet), context); - final boolean preserveResponseBodyAsStream = HttpUtils.shouldPreserveResponseBodyAsStream(context); - final Object result; if (TypeUtil.isTypeOrSubTypeOf(returnType, Mono.class)) { final Type monoTypeParam = TypeUtil.getTypeArgument(returnType); @@ -266,8 +261,8 @@ private Object handleRestReturnType(Mono> - result = asyncExpectedResponse.flatMap(response -> handleRestResponseReturnType(response, methodParser, - monoTypeParam, preserveResponseBodyAsStream)); + result = asyncExpectedResponse + .flatMap(response -> handleRestResponseReturnType(response, methodParser, monoTypeParam)); } } else if (FluxUtil.isFluxByteBuffer(returnType)) { // ProxyMethod ReturnType: Flux @@ -280,8 +275,9 @@ private Object handleRestReturnType(Mono handleRestResponseReturnType(httpResponse, - methodParser, returnType, preserveResponseBodyAsStream)).block(); + result = asyncExpectedResponse + .flatMap(httpResponse -> handleRestResponseReturnType(httpResponse, methodParser, returnType)) + .block(); } return result; } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java index cd9441bb5548..b1b370f8757f 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java @@ -153,14 +153,6 @@ protected abstract Object invoke(Object proxy, Method method, RequestOptions opt EnumSet errorOptions, Consumer httpRequestConsumer, SwaggerMethodParser methodParser, HttpRequest request, Context context); - final Context updateRequestContext(HttpRequest request, Context context) { - if (HttpUtils.acceptsTextEventStream(request.getHeaders().getValue(HttpHeaderName.ACCEPT))) { - return context.addData(HttpUtils.AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM, true); - } - - return context; - } - /** * Update the request with the provided configuration. * diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java index f4269a0e0a4b..e1f8890317a1 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java @@ -78,8 +78,6 @@ public Object invoke(Object proxy, Method method, RequestOptions options, EnumSe requestCallback.accept(request); } - context = updateRequestContext(request, context); - if (request.getBodyAsBinaryData() != null) { request.setBody(RestProxyUtils.validateLengthSync(request)); } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java index bf9f41a7512e..4c14fc3e57c1 100644 --- a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java +++ b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java @@ -3,7 +3,6 @@ package com.azure.core.implementation.util; import com.azure.core.util.Configuration; -import com.azure.core.util.Context; import com.azure.core.util.logging.ClientLogger; import java.time.Duration; @@ -48,12 +47,6 @@ public final class HttpUtils { */ public static final String AZURE_EAGERLY_READ_RESPONSE = "azure-eagerly-read-response"; - /** - * Context key that instructs REST proxy response ownership and decoding to preserve the response body as a live - * stream. HTTP client implementations do not consume this key. - */ - public static final String AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM = "azure-preserve-response-body-as-stream"; - /** * Context key used to indicate to an HttpClient implementation if the response body should be ignored and eagerly * drained from the network. @@ -71,36 +64,6 @@ public final class HttpUtils { */ public static final String AZURE_EAGERLY_CONVERT_HEADERS = "azure-eagerly-convert-headers"; - /** - * Determines whether the response body must be preserved as a live stream. - * - * @param context Contextual information about the request. - * @return Whether the response body must be preserved as a live stream. - */ - public static boolean shouldPreserveResponseBodyAsStream(Context context) { - return Boolean.TRUE.equals(context.getData(AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM).orElse(false)); - } - - /** - * Determines whether an Accept header contains a {@code text/event-stream} media range. - * - * @param headerValue The header value. - * @return Whether the header contains a {@code text/event-stream} media range. - */ - public static boolean acceptsTextEventStream(String headerValue) { - if (headerValue == null) { - return false; - } - - for (String value : splitHeaderValue(headerValue, ',')) { - if (isTextEventStreamContentType(value)) { - return true; - } - } - - return false; - } - /** * Determines whether a Content-Type header identifies exactly one {@code text/event-stream} representation. * Charset parameters don't affect this determination as event streams are always decoded as UTF-8. diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java index 3db53cbc3b66..cdad461025c5 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/HttpLoggingPolicyTests.java @@ -16,7 +16,6 @@ import com.azure.core.implementation.AccessibleByteArrayOutputStream; import com.azure.core.implementation.accesshelpers.ClientLoggerAccessHelper; import com.azure.core.implementation.logging.DefaultLogger; -import com.azure.core.implementation.util.HttpUtils; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.core.util.FluxUtil; @@ -275,8 +274,8 @@ public void validateLoggingDoesNotConsumeResponseSync(BinaryData responseBody, b @ParameterizedTest(name = "[{index}] {displayName}") @MethodSource("responseLoggingSupplier") - public void responseLoggingUsesActualContentType(String contentType, boolean useStreamingContext, - int expectedBufferCount, boolean expectBodyLogged) { + public void responseLoggingUsesActualContentType(String contentType, int expectedBufferCount, + boolean expectBodyLogged) { byte[] data = "streaming response".getBytes(StandardCharsets.UTF_8); AtomicInteger bufferCount = new AtomicInteger(); HttpRequest request = new HttpRequest(HttpMethod.GET, "https://test.com/responseLoggingUsesActualContentType"); @@ -290,9 +289,6 @@ public void responseLoggingUsesActualContentType(String contentType, boolean use .build(); Context context = getCallerMethodContext("streamingResponsesAreNotBuffered", LogLevel.INFORMATIONAL); - if (useStreamingContext) { - context = context.addData(HttpUtils.AZURE_PRESERVE_RESPONSE_BODY_AS_STREAM, true); - } try (HttpResponse response = pipeline.send(request, context).block()) { assertNotNull(response); @@ -309,8 +305,8 @@ public void responseLoggingUsesActualContentType(String contentType, boolean use } private static Stream responseLoggingSupplier() { - return Stream.of(Arguments.of("Text/Event-Stream; charset=utf-8", false, 0, false), - Arguments.of(ContentType.APPLICATION_JSON, true, 2, true)); + return Stream.of(Arguments.of("Text/Event-Stream; charset=utf-8", 0, false), + Arguments.of(ContentType.APPLICATION_JSON, 2, true)); } private static Stream validateLoggingDoesNotConsumeSupplierSync() { diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java index 1b5244a72e6c..ed366c1be3c2 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/http/rest/AsyncRestProxyTests.java @@ -39,7 +39,6 @@ import static com.azure.core.CoreTestUtils.assertArraysEqual; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -210,7 +209,8 @@ public static Stream getResponseHeaderAndReplayability() { @ParameterizedTest @MethodSource("streamingResponseOwnershipSupplier") - public void streamingResponseIsUnbufferedAndClosedOnConsumption(String accept, String contentType) { + public void streamingResponseUsesResponseContentType(String accept, String contentType, boolean replayable, + int expectedCloseCount) { byte[] expectedBytes = "hello".getBytes(StandardCharsets.UTF_8); AtomicInteger responseCloseCount = new AtomicInteger(); HttpClient client = request -> Mono.just(new MockHttpResponse(request, 200, @@ -229,9 +229,9 @@ public void close() { Response response = service.getStreamingResponse(options, Context.NONE).block(); - assertFalse(response.getValue().isReplayable()); + assertEquals(replayable, response.getValue().isReplayable()); assertArraysEqual(expectedBytes, response.getValue().toBytes()); - assertEquals(1, responseCloseCount.get()); + assertEquals(expectedCloseCount, responseCloseCount.get()); } @Test @@ -264,7 +264,7 @@ public void close() { } private static Stream streamingResponseOwnershipSupplier() { - return Stream.of(Arguments.of("text/event-stream", ContentType.APPLICATION_JSON), - Arguments.of(null, "text/event-stream; charset=utf-8")); + return Stream.of(Arguments.of("text/event-stream", ContentType.APPLICATION_JSON, true, 0), + Arguments.of(null, "text/event-stream; charset=utf-8", false, 1)); } } diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java index 82448e7eeb80..430a3772e8fe 100644 --- a/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java +++ b/sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/HttpUtilsTests.java @@ -9,14 +9,6 @@ import static org.junit.jupiter.api.Assertions.assertTrue; public class HttpUtilsTests { - @Test - public void acceptsTextEventStreamIgnoresParameters() { - assertTrue(HttpUtils.acceptsTextEventStream("text/event-stream;q=0")); - assertTrue(HttpUtils.acceptsTextEventStream("application/json, text/event-stream;q=invalid")); - assertFalse(HttpUtils.acceptsTextEventStream("application/json; note=\"text/event-stream, q=1\"")); - assertTrue(HttpUtils.acceptsTextEventStream("text/event-stream; note=\"x,y;q=0.5\"")); - } - @Test public void textEventStreamContentTypeRequiresSingleMediaType() { assertTrue(HttpUtils.isTextEventStreamContentType("Text/Event-Stream; charset=utf-8"));