Add BinaryData.isReplayable (Azure#29463)

* Add `BinaryData.isReplayable`

* pr feedback.

Author: kasobol-msft

sdk/core/azure-core/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 
+- Added `BinaryData.isReplayable()` to indicate if multiple consumptions of the content are safe.
 - Added support for sending synchronous requests using `sendSync` in `HttpPipeline`:
   - Added `HttpPipelinePolicy.processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next)` to allow processing policies synchronously.
   - Added `HttpPipelineSyncPolicy` to represent synchronous `HttpPipelinePolicy`.

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/BinaryDataContent.java

@@ -74,4 +74,11 @@ public abstract class BinaryDataContent {
      * @return The {@link BinaryDataContent} as a {@code Flux<ByteBuffer>}.
      */
     public abstract Flux<ByteBuffer> toFluxByteBuffer();
+
+    /**
+     * Returns a flag indicating whether the content can be repeatedly consumed using all accessors including
+     * {@link #toStream()} and {@link #toFluxByteBuffer()}
+     * @return a flag indicating whether the content can be repeatedly consumed.
+     */
+    public abstract boolean isReplayable();
 }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/ByteArrayContent.java

@@ -64,4 +64,9 @@ public ByteBuffer toByteBuffer() {
     public Flux<ByteBuffer> toFluxByteBuffer() {
         return Mono.fromSupplier(() -> ByteBuffer.wrap(toBytes()).asReadOnlyBuffer()).flux();
     }
+
+    @Override
+    public boolean isReplayable() {
+        return true;
+    }
 }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FileContent.java

@@ -195,6 +195,11 @@ public int getChunkSize() {
         return chunkSize;
     }
 
+    @Override
+    public boolean isReplayable() {
+        return true;
+    }
+
     private byte[] getBytes() {
         if (length > MAX_ARRAY_SIZE) {
             throw LOGGER.logExceptionAsError(new IllegalStateException(

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java

@@ -87,6 +87,11 @@ public Flux<ByteBuffer> toFluxByteBuffer() {
         return content;
     }
 
+    @Override
+    public boolean isReplayable() {
+        return false;
+    }
+
     private byte[] getBytes() {
         return FluxUtil.collectBytesInByteBufferStream(content)
                 // this doesn't seem to be working (newBoundedElastic() didn't work either)

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/InputStreamContent.java

@@ -80,6 +80,11 @@ public Flux<ByteBuffer> toFluxByteBuffer() {
         return FluxUtil.toFluxByteBuffer(this.content, STREAM_READ_SIZE);
     }
 
+    @Override
+    public boolean isReplayable() {
+        return false;
+    }
+
     private byte[] getBytes() {
         try {
             ByteArrayOutputStream dataOutputBuffer = new ByteArrayOutputStream();

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/SerializableContent.java

@@ -76,6 +76,11 @@ public Flux<ByteBuffer> toFluxByteBuffer() {
         return Mono.fromSupplier(() -> ByteBuffer.wrap(toBytes())).flux();
     }
 
+    @Override
+    public boolean isReplayable() {
+        return true;
+    }
+
     private byte[] getBytes() {
         return serializer.serializeToBytes(content);
     }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/StringContent.java

@@ -71,6 +71,11 @@ public Flux<ByteBuffer> toFluxByteBuffer() {
         return Mono.fromSupplier(() -> ByteBuffer.wrap(toBytes())).flux();
     }
 
+    @Override
+    public boolean isReplayable() {
+        return true;
+    }
+
     private byte[] getBytes() {
         return this.content.getBytes(StandardCharsets.UTF_8);
     }

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

@@ -1459,4 +1459,13 @@ public Flux<ByteBuffer> toFluxByteBuffer() {
     public Long getLength() {
         return content.getLength();
     }
+
+    /**
+     * Returns a flag indicating whether the content can be repeatedly consumed using all accessors including
+     * {@link #toStream()} and {@link #toFluxByteBuffer()}
+     * @return a flag indicating whether the content can be repeatedly consumed using all accessors.
+     */
+    public boolean isReplayable() {
+        return content.isReplayable();
+    }
 }

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

@@ -10,8 +10,11 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.type.TypeFactory;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Named;
 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 org.junit.jupiter.params.provider.ValueSource;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -44,13 +47,15 @@
 import java.util.Random;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 
 import static com.azure.core.implementation.util.BinaryDataContent.STREAM_READ_SIZE;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doThrow;
@@ -622,6 +627,113 @@ public void testFromFileSegment(int size) throws Exception {
 
     }
 
+    @ParameterizedTest
+    @MethodSource("createNonRetryableBinaryData")
+    public void testNonReplayableContentTypes(Supplier<BinaryData> binaryDataSupplier) throws IOException {
+
+        assertFalse(binaryDataSupplier.get().isReplayable());
+
+        BinaryData data = binaryDataSupplier.get();
+        byte[] firstFluxConsumption = FluxUtil.collectBytesInByteBufferStream(data.toFluxByteBuffer()).block();
+        byte[] secondFluxConsumption = FluxUtil.collectBytesInByteBufferStream(data.toFluxByteBuffer()).block();
+
+        data = binaryDataSupplier.get();
+        byte[] firstStreamConsumption = readInputStream(data.toStream());
+        byte[] secondStreamConsumption = readInputStream(data.toStream());
+
+        // Either flux or stream consumption is not replayable.
+        assertFalse(
+            Arrays.equals(firstFluxConsumption, secondFluxConsumption)
+            && Arrays.equals(firstStreamConsumption, secondStreamConsumption)
+        );
+    }
+
+    public static Stream<Arguments> createNonRetryableBinaryData() {
+        byte[] bytes = new byte[1024];
+        RANDOM.nextBytes(bytes);
+        return Stream.of(
+            Arguments.of(
+                Named.named("stream",
+                    (Supplier<BinaryData>) () -> BinaryData.fromStream(new ByteArrayInputStream(bytes)))),
+            Arguments.of(
+                Named.named("unbuffered flux",
+                    (Supplier<BinaryData>) () -> BinaryData.fromFlux(Flux.just(ByteBuffer.wrap(bytes)), null, false).block()))
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("createRetryableBinaryData")
+    public void testReplayableContentTypes(Supplier<BinaryData> binaryDataSupplier) throws IOException {
+
+        assertTrue(binaryDataSupplier.get().isReplayable());
+
+        // Check toFluxByteBuffer consumption
+        BinaryData data = binaryDataSupplier.get();
+        byte[] firstConsumption = FluxUtil.collectBytesInByteBufferStream(data.toFluxByteBuffer()).block();
+        byte[] secondConsumption = FluxUtil.collectBytesInByteBufferStream(data.toFluxByteBuffer()).block();
+        assertArrayEquals(firstConsumption, secondConsumption);
+
+        // Check toStream consumption
+        data = binaryDataSupplier.get();
+        firstConsumption = readInputStream(data.toStream());
+        secondConsumption = readInputStream(data.toStream());
+        assertArrayEquals(firstConsumption, secondConsumption);
+
+        // Check toByteBuffer consumption
+        data = binaryDataSupplier.get();
+        firstConsumption = readByteBuffer(data.toByteBuffer());
+        secondConsumption = readByteBuffer(data.toByteBuffer());
+        assertArrayEquals(firstConsumption, secondConsumption);
+
+        // Check toBytes consumption
+        data = binaryDataSupplier.get();
+        firstConsumption = data.toBytes();
+        secondConsumption = data.toBytes();
+        assertArrayEquals(firstConsumption, secondConsumption);
+    }
+
+    public static Stream<Arguments> createRetryableBinaryData() throws IOException {
+        byte[] bytes = new byte[1024];
+        RANDOM.nextBytes(bytes);
+        Path tempFile = Files.createTempFile("retryableData", null);
+        tempFile.toFile().deleteOnExit();
+        Files.write(tempFile, bytes);
+        return Stream.of(
+            Arguments.of(
+                Named.named("bytes",
+                    (Supplier<BinaryData>) () -> BinaryData.fromBytes(bytes))),
+            Arguments.of(
+                Named.named("string",
+                    (Supplier<BinaryData>) () -> BinaryData.fromString("test string"))),
+            Arguments.of(
+                Named.named("object",
+                    (Supplier<BinaryData>) () -> BinaryData.fromObject("\"test string\""))),
+            Arguments.of(
+                Named.named("file",
+                    (Supplier<BinaryData>) () -> BinaryData.fromFile(tempFile))),
+            Arguments.of(
+                Named.named("buffered flux",
+                    (Supplier<BinaryData>) () -> BinaryData.fromFlux(Flux.just(ByteBuffer.wrap(bytes))).block()))
+        );
+    }
+
+    private static byte[] readInputStream(InputStream inputStream) throws IOException {
+        byte[] buffer = new byte[1024];
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        int read;
+        while ((read = inputStream.read(buffer)) >= 0) {
+            bos.write(buffer, 0, read);
+        }
+        return bos.toByteArray();
+    }
+
+    private static byte[] readByteBuffer(ByteBuffer buffer) {
+        // simplified implementation good enough for testing.
+        byte[] result = new byte[buffer.remaining()];
+        buffer.get(result);
+        return result;
+    }
+
     public static class MyJsonSerializer implements JsonSerializer {
         private final ClientLogger logger = new ClientLogger(MyJsonSerializer.class);
         private final ObjectMapper mapper;