Reduce Buffering in SDKs (Azure#26821)

Reduce Buffering in SDKs

Author: alzimmermsft

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -396,6 +396,7 @@
       <Class name="com.azure.core.util.FluxUtil"/>
       <Class name="com.azure.core.util.FluxUtil$1"/>
       <Class name="com.azure.core.implementation.FileWriteSubscriber"/>
+      <Class name="com.azure.core.implementation.OutputStreamWriteSubscriber"/>
     </Or>
     <Bug pattern="UWF_FIELD_NOT_INITIALIZED_IN_CONSTRUCTOR"/>
   </Match>

eng/versioning/version_client.txt

@@ -345,6 +345,8 @@ com.azure.tools:azure-sdk-archetype;1.0.0;1.0.0
 # In the pom, the version update tag after the version should name the unreleased package and the dependency version:
 # <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
 
+unreleased_com.azure:azure-core;1.26.0-beta.1
+
 # Released Beta dependencies: Copy the entry from above, prepend "beta_", remove the current
 # version and set the version to the released beta. Released beta dependencies are only valid
 # for dependency versions. These entries are specifically for when we've released a beta for

sdk/communication/azure-communication-callingserver/pom.xml

@@ -53,7 +53,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core</artifactId>
-      <version>1.25.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
+      <version>1.26.0-beta.1</version> <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>

sdk/communication/azure-communication-callingserver/src/main/java/com/azure/communication/callingserver/ContentDownloader.java

@@ -15,7 +15,6 @@
 import com.azure.core.util.Context;
 import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
-import reactor.core.Exceptions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.SignalType;
@@ -54,15 +53,9 @@ Mono<Response<Void>> downloadToStreamWithResponse(
         HttpRange httpRange,
         Context context) {
         return downloadStreamWithResponse(sourceEndpoint, httpRange, context)
-            .flatMap(response -> response.getValue().reduce(destinationStream, (outputStream, buffer) -> {
-                try {
-                    outputStream.write(FluxUtil.byteBufferToArray(buffer));
-                    return outputStream;
-                } catch (IOException ex) {
-                    throw logger.logExceptionAsError(Exceptions.propagate(new UncheckedIOException(ex)));
-                }
-            }).thenReturn(new SimpleResponse<>(response.getRequest(), response.getStatusCode(),
-                response.getHeaders(), null)));
+            .flatMap(response -> FluxUtil.writeToOutputStream(response.getValue(), destinationStream)
+                .thenReturn(new SimpleResponse<>(response.getRequest(), response.getStatusCode(),
+                    response.getHeaders(), null)));
     }
 
     Mono<Response<Flux<ByteBuffer>>> downloadStreamWithResponse(

sdk/communication/azure-communication-callingserver/src/test/java/com/azure/communication/callingserver/DownloadContentLiveTests.java

@@ -13,6 +13,7 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.io.UncheckedIOException;
 import java.io.UnsupportedEncodingException;
 import java.nio.charset.StandardCharsets;
@@ -130,12 +131,12 @@ public void downloadContentStreamFailure(HttpClient httpClient) throws IOExcepti
         CallingServerClientBuilder builder = getCallingServerClientUsingConnectionString(httpClient);
         CallingServerClient conversationClient = setupClient(builder, "downloadContent404");
 
-        ByteArrayOutputStream byteArrayOutputStream = Mockito.mock(ByteArrayOutputStream.class);
-        doThrow(IOException.class).when(byteArrayOutputStream).write(Mockito.any());
+        OutputStream outputStream = Mockito.mock(OutputStream.class);
+        doThrow(IOException.class).when(outputStream).write(Mockito.any(), Mockito.anyInt(), Mockito.anyInt());
         assertThrows(
             UncheckedIOException.class,
             () -> conversationClient
-                .downloadTo(METADATA_URL, byteArrayOutputStream, null));
+                .downloadTo(METADATA_URL, outputStream, null));
     }
 
     private CallingServerClient setupClient(CallingServerClientBuilder builder, String testName) {

sdk/core/azure-core/src/main/java/com/azure/core/implementation/ImplUtils.java

@@ -6,7 +6,12 @@
 import com.azure.core.http.HttpHeaders;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.DateTimeRfc1123;
+import com.azure.core.util.FluxUtil;
 
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.time.DateTimeException;
 import java.time.Duration;
 import java.time.OffsetDateTime;
@@ -94,6 +99,47 @@ private static long tryParseLong(String value) {
         }
     }
 
+    /**
+     * Writes a {@link ByteBuffer} into an {@link OutputStream}.
+     * <p>
+     * This method provides writing optimization based on the type of {@link ByteBuffer} and {@link OutputStream}
+     * passed. For example, if the {@link ByteBuffer} has a backing {@code byte[]} this method will access that directly
+     * to write to the {@code stream} instead of buffering the contents of the {@link ByteBuffer} into a temporary
+     * buffer.
+     *
+     * @param buffer The {@link ByteBuffer} to write into the {@code stream}.
+     * @param stream The {@link OutputStream} where the {@code buffer} will be written.
+     * @throws IOException If an I/O occurs while writing the {@code buffer} into the {@code stream}.
+     */
+    public static void writeByteBufferToStream(ByteBuffer buffer, OutputStream stream) throws IOException {
+        // First check if the buffer has a backing byte[]. The backing byte[] can be accessed directly and written
+        // without an additional buffering byte[].
+        if (buffer.hasArray()) {
+            // Write the byte[] from the current view position to the length remaining in the view.
+            stream.write(buffer.array(), buffer.position(), buffer.remaining());
+
+            // Update the position of the ByteBuffer to treat this the same as getting from the buffer.
+            buffer.position(buffer.position() + buffer.remaining());
+            return;
+        }
+
+        // Next begin checking for specific instances of OutputStream that may provide better writing options for
+        // direct ByteBuffers.
+        if (stream instanceof FileOutputStream) {
+            FileOutputStream fileOutputStream = (FileOutputStream) stream;
+
+            // Writing to the FileChannel directly may provide native optimizations for moving the OS managed memory
+            // into the file.
+            // Write will move both the OutputStream's and ByteBuffer's position so there is no need to perform
+            // additional updates that are required when using the backing array.
+            fileOutputStream.getChannel().write(buffer);
+            return;
+        }
+
+        // All optimizations have been exhausted, fallback to buffering write.
+        stream.write(FluxUtil.byteBufferToArray(buffer));
+    }
+
     private ImplUtils() {
     }
 }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/OutputStreamWriteSubscriber.java

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.core.implementation;
+
+import com.azure.core.util.logging.ClientLogger;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.MonoSink;
+import reactor.core.publisher.Operators;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+
+/**
+ * Subscriber that writes a stream of {@link ByteBuffer ByteBuffers} to an {@link OutputStream}.
+ */
+@SuppressWarnings("ReactiveStreamsSubscriberImplementation")
+public final class OutputStreamWriteSubscriber implements Subscriber<ByteBuffer> {
+    private final MonoSink<Void> emitter;
+    private final OutputStream stream;
+    private final ClientLogger logger;
+
+    private Subscription subscription;
+
+    public OutputStreamWriteSubscriber(MonoSink<Void> emitter, OutputStream stream, ClientLogger logger) {
+        this.emitter = emitter;
+        this.stream = stream;
+        this.logger = logger;
+    }
+
+    @Override
+    public void onSubscribe(Subscription s) {
+        // Only set the Subscription if one has not been previously set.
+        // Any additional Subscriptions will be cancelled.
+        if (Operators.validate(this.subscription, s)) {
+            subscription = s;
+
+            s.request(1);
+        }
+    }
+
+    @Override
+    public void onNext(ByteBuffer byteBuffer) {
+        try {
+            ImplUtils.writeByteBufferToStream(byteBuffer, stream);
+            subscription.request(1);
+        } catch (IOException ex) {
+            // Emit IOExceptions as UncheckIOExceptions as that is the pattern used in SDKs.
+            onError(new UncheckedIOException(ex));
+        }
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+        subscription.cancel();
+        emitter.error(logger.logThrowableAsError(throwable));
+    }
+
+    @Override
+    public void onComplete() {
+        emitter.success();
+    }
+}

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

@@ -8,6 +8,7 @@
 import com.azure.core.http.rest.Response;
 import com.azure.core.implementation.ByteBufferCollector;
 import com.azure.core.implementation.FileWriteSubscriber;
+import com.azure.core.implementation.OutputStreamWriteSubscriber;
 import com.azure.core.implementation.RetriableDownloadFlux;
 import com.azure.core.implementation.TypeUtil;
 import com.azure.core.util.logging.ClientLogger;
@@ -21,12 +22,15 @@
 import reactor.core.publisher.Operators;
 import reactor.util.context.ContextView;
 
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.lang.reflect.Type;
 import java.nio.ByteBuffer;
 import java.nio.channels.AsynchronousFileChannel;
 import java.nio.channels.CompletionHandler;
+import java.nio.channels.FileChannel;
 import java.nio.file.StandardOpenOption;
 import java.util.Collections;
 import java.util.Map;
@@ -42,6 +46,7 @@
  */
 public final class FluxUtil {
     private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
+    private static final ClientLogger LOGGER = new ClientLogger(FluxUtil.class);
 
     /**
      * Checks if a type is Flux<ByteBuffer>.
@@ -203,6 +208,40 @@ public static Flux<ByteBuffer> toFluxByteBuffer(InputStream inputStream, int chu
             return Flux.empty();
         }
 
+        // If the InputStream is an instance of FileInputStream we should be able to leverage the FileChannel backing
+        // the FileInputStream to generated MappedByteBuffers which aren't loaded into memory until the content is
+        // consumed. This at least defers the memory usage until later and may also provide downstream calls ways to
+        // optimize if they have special cases for MappedByteBuffer.
+        if (inputStream instanceof FileInputStream) {
+            FileChannel fileChannel = ((FileInputStream) inputStream).getChannel();
+
+            return Flux.<ByteBuffer, FileChannel>generate(() -> fileChannel, (channel, sink) -> {
+                try {
+                    long channelPosition = channel.position();
+                    long channelSize = channel.size();
+
+                    if (channelPosition == channelSize) {
+                        // End of File has been reached, signal completion.
+                        channel.close();
+                        sink.complete();
+                    } else {
+                        // Determine the size of the next MappedByteBuffer, either the remaining File contents or the
+                        // expected chunk size.
+                        int nextByteBufferSize = (int) Math.min(chunkSize, channelSize - channelPosition);
+                        sink.next(channel.map(FileChannel.MapMode.READ_ONLY, channelPosition, nextByteBufferSize));
+
+                        // FileChannel.map doesn't update the FileChannel's position as reading would, so the position
+                        // needs to be updated based on the number of bytes mapped.
+                        channel.position(channelPosition + nextByteBufferSize);
+                    }
+                } catch (IOException ex) {
+                    sink.error(ex);
+                }
+
+                return channel;
+            });
+        }
+
         return Flux.<ByteBuffer, InputStream>generate(() -> inputStream, (stream, sink) -> {
             byte[] buffer = new byte[chunkSize];
 
@@ -420,6 +459,32 @@ public static reactor.util.context.Context toReactorContext(Context context) {
         return returnContext;
     }
 
+    /**
+     * Writes the {@link ByteBuffer ByteBuffers} emitted by a {@link Flux} of {@link ByteBuffer} to an {@link
+     * OutputStream}.
+     * <p>
+     * The {@code stream} is not closed by this call, closing of the {@code stream} is managed by the caller.
+     * <p>
+     * The response {@link Mono} will emit an error if {@code content} or {@code stream} are null. Additionally, an
+     * error will be emitted if an exception occurs while writing the {@code content} to the {@code stream}.
+     *
+     * @param content The {@link Flux} of {@link ByteBuffer} content.
+     * @param stream The {@link OutputStream} being written into.
+     * @return A {@link Mono} which emits a completion status once the {@link Flux} has been written to the {@link
+     * OutputStream}, or an error status if writing fails.
+     */
+    public static Mono<Void> writeToOutputStream(Flux<ByteBuffer> content, OutputStream stream) {
+        if (content == null && stream == null) {
+            return monoError(LOGGER, new NullPointerException("'content' and 'stream' cannot be null."));
+        } else if (content == null) {
+            return monoError(LOGGER, new NullPointerException("'content' cannot be null."));
+        } else if (stream == null) {
+            return monoError(LOGGER, new NullPointerException("'stream' cannot be null."));
+        }
+
+        return Mono.create(emitter -> content.subscribe(new OutputStreamWriteSubscriber(emitter, stream, LOGGER)));
+    }
+
     /**
      * Writes the {@link ByteBuffer ByteBuffers} emitted by a {@link Flux} of {@link ByteBuffer} to an {@link
      * AsynchronousFileChannel}.
@@ -456,24 +521,17 @@ public static Mono<Void> writeFile(Flux<ByteBuffer> content, AsynchronousFileCha
      * AsynchronousFileChannel}.
      */
     public static Mono<Void> writeFile(Flux<ByteBuffer> content, AsynchronousFileChannel outFile, long position) {
-        return Mono.create(emitter -> {
-            if (content == null) {
-                emitter.error(new NullPointerException("'content' cannot be null."));
-                return;
-            }
-
-            if (outFile == null) {
-                emitter.error(new NullPointerException("'outFile' cannot be null."));
-                return;
-            }
-
-            if (position < 0) {
-                emitter.error(new IllegalArgumentException("'position' cannot be less than 0."));
-                return;
-            }
+        if (content == null && outFile == null) {
+            return monoError(LOGGER, new NullPointerException("'content' and 'outFile' cannot be null."));
+        } else if (content == null) {
+            return monoError(LOGGER, new NullPointerException("'content' cannot be null."));
+        } else if (outFile == null) {
+            return monoError(LOGGER, new NullPointerException("'outFile' cannot be null."));
+        } else if (position < 0) {
+            return monoError(LOGGER, new IllegalArgumentException("'position' cannot be less than 0."));
+        }
 
-            content.subscribe(new FileWriteSubscriber(outFile, position, emitter));
-        });
+        return Mono.create(emitter -> content.subscribe(new FileWriteSubscriber(outFile, position, emitter)));
     }
 
     /**