From 19cc55c8a7964b167226e8edfedd17410e843cde Mon Sep 17 00:00:00 2001 From: Isabelle Date: Wed, 10 Dec 2025 10:11:32 -0800 Subject: [PATCH 1/4] adjusting encoder logic --- .../StructuredMessageEncoder.java | 128 ++++++++---------- 1 file changed, 55 insertions(+), 73 deletions(-) diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java index ca02920a13f0..365773b09631 100644 --- a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java +++ b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java @@ -6,12 +6,13 @@ import com.azure.core.util.logging.ClientLogger; import com.azure.storage.common.implementation.StorageCrc64Calculator; import com.azure.storage.common.implementation.StorageImplUtils; +import reactor.core.publisher.Flux; -import java.io.IOException; import java.nio.ByteBuffer; -import java.io.ByteArrayOutputStream; import java.nio.ByteOrder; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.CRC64_LENGTH; @@ -35,7 +36,6 @@ public class StructuredMessageEncoder { private int currentContentOffset; private int currentSegmentNumber; private int currentSegmentOffset; - private int currentMessageLength; private long messageCRC64; private final Map segmentCRC64s; @@ -66,7 +66,6 @@ public StructuredMessageEncoder(int contentLength, int segmentSize, StructuredMe this.currentSegmentOffset = 0; this.messageCRC64 = 0; this.segmentCRC64s = new HashMap<>(); - this.currentMessageLength = 0; if (numSegments > Short.MAX_VALUE) { StorageImplUtils.assertInBounds("numSegments", numSegments, 1, Short.MAX_VALUE); @@ -109,115 +108,99 @@ private byte[] generateMessageHeader() { return buffer.array(); } - private byte[] generateSegmentHeader() { - int segmentHeaderSize = Math.min(segmentSize, contentLength - currentContentOffset); - // 2 byte number, 8 byte size + private byte[] generateSegmentHeader(int segmentContentSize) { ByteBuffer buffer = ByteBuffer.allocate(getSegmentHeaderLength()).order(ByteOrder.LITTLE_ENDIAN); buffer.putShort((short) currentSegmentNumber); - buffer.putLong(segmentHeaderSize); + buffer.putLong(segmentContentSize); return buffer.array(); } /** - * Encodes the given buffer into a structured message format. + * Encodes the given buffer into a structured message format as a stream of ByteBuffers. * * @param unencodedBuffer The buffer to be encoded. - * @return The encoded buffer. - * @throws IOException If an error occurs while encoding the buffer. + * @return A Flux of encoded ByteBuffers. * @throws IllegalArgumentException If the buffer length exceeds the content length, or the content has already been * encoded. */ - public ByteBuffer encode(ByteBuffer unencodedBuffer) throws IOException { + public Flux encode(ByteBuffer unencodedBuffer) { StorageImplUtils.assertNotNull("unencodedBuffer", unencodedBuffer); if (currentContentOffset == contentLength) { - throw LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded.")); + return Flux + .error(LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded."))); } if ((unencodedBuffer.remaining() + currentContentOffset) > contentLength) { - throw LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length.")); + return Flux.error( + LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length."))); } if (!unencodedBuffer.hasRemaining()) { - return ByteBuffer.allocate(0); + return Flux.empty(); } - ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + return Flux.defer(() -> { + List buffers = new ArrayList<>(); - // if we are at the beginning of the message, encode message header - if (currentMessageLength == 0) { - encodeMessageHeader(byteArrayOutputStream); - } - - while (unencodedBuffer.hasRemaining()) { - // if we are at the beginning of a segment's content, encode segment header - if (currentSegmentOffset == 0) { - encodeSegmentHeader(byteArrayOutputStream); + // if we are at the beginning of the message, encode message header + if (currentContentOffset == 0) { + buffers.add(ByteBuffer.wrap(generateMessageHeader())); } - encodeSegmentContent(unencodedBuffer, byteArrayOutputStream); - - // if we are at the end of a segment's content, encode segment footer - if (currentSegmentOffset == getSegmentContentLength()) { - encodeSegmentFooter(byteArrayOutputStream); + while (unencodedBuffer.hasRemaining()) { + // if we are at the beginning of a segment's content, encode segment header + if (currentSegmentOffset == 0) { + incrementCurrentSegment(); + // Calculate actual segment size based on remaining content + int actualSegmentSize = Math.min(segmentSize, contentLength - currentContentOffset); + buffers.add(ByteBuffer.wrap(generateSegmentHeader(actualSegmentSize))); + } + + buffers.add(encodeSegmentContent(unencodedBuffer)); + + // if we are at the end of a segment's content, encode segment footer + if (currentSegmentOffset == getSegmentContentLength()) { + byte[] footer = generateSegmentFooter(); + if (footer.length > 0) { + buffers.add(ByteBuffer.wrap(footer)); + } + currentSegmentOffset = 0; + } } - } - - // if all content has been encoded, encode message footer - if (currentContentOffset == contentLength) { - encodeMessageFooter(byteArrayOutputStream); - } - - return ByteBuffer.wrap(byteArrayOutputStream.toByteArray()); - } - private void encodeMessageHeader(ByteArrayOutputStream output) { - byte[] metadata = generateMessageHeader(); - output.write(metadata, 0, metadata.length); - - currentMessageLength += metadata.length; - } - - private void encodeSegmentHeader(ByteArrayOutputStream output) { - incrementCurrentSegment(); - byte[] metadata = generateSegmentHeader(); - output.write(metadata, 0, metadata.length); + // if all content has been encoded, encode message footer + if (currentContentOffset == contentLength) { + byte[] footer = generateMessageFooter(); + if (footer.length > 0) { + buffers.add(ByteBuffer.wrap(footer)); + } + } - currentMessageLength += metadata.length; + return Flux.fromIterable(buffers); + }); } - private void encodeSegmentFooter(ByteArrayOutputStream output) { - byte[] metadata; + private byte[] generateSegmentFooter() { if (structuredMessageFlags == StructuredMessageFlags.STORAGE_CRC64) { - metadata = ByteBuffer.allocate(CRC64_LENGTH) + return ByteBuffer.allocate(CRC64_LENGTH) .order(ByteOrder.LITTLE_ENDIAN) .putLong(segmentCRC64s.get(currentSegmentNumber)) .array(); - } else { - metadata = new byte[0]; } - output.write(metadata, 0, metadata.length); - - currentMessageLength += metadata.length; - currentSegmentOffset = 0; + return new byte[0]; } - private void encodeMessageFooter(ByteArrayOutputStream output) { - byte[] metadata; + private byte[] generateMessageFooter() { if (structuredMessageFlags == StructuredMessageFlags.STORAGE_CRC64) { - metadata = ByteBuffer.allocate(CRC64_LENGTH).order(ByteOrder.LITTLE_ENDIAN).putLong(messageCRC64).array(); - } else { - metadata = new byte[0]; + return ByteBuffer.allocate(CRC64_LENGTH).order(ByteOrder.LITTLE_ENDIAN).putLong(messageCRC64).array(); } - - output.write(metadata, 0, metadata.length); - currentMessageLength += metadata.length; + return new byte[0]; } - private void encodeSegmentContent(ByteBuffer unencodedBuffer, ByteArrayOutputStream output) { + private ByteBuffer encodeSegmentContent(ByteBuffer unencodedBuffer) { int readSize = Math.min(unencodedBuffer.remaining(), getSegmentContentLength() - currentSegmentOffset); - byte[] content = new byte[readSize]; unencodedBuffer.get(content, 0, readSize); @@ -230,8 +213,7 @@ private void encodeSegmentContent(ByteBuffer unencodedBuffer, ByteArrayOutputStr currentContentOffset += readSize; currentSegmentOffset += readSize; - output.write(content, 0, content.length); - currentMessageLength += readSize; + return ByteBuffer.wrap(content); } private int calculateMessageLength() { @@ -255,7 +237,7 @@ private void incrementCurrentSegment() { * * @return The length of the message. */ - public int getMessageLength() { + public long getEncodedMessageLength() { return messageLength; } } From 753f5b58b01e33f0a3f6d79113f7ff5bbe932d9b Mon Sep 17 00:00:00 2001 From: Isabelle Date: Wed, 10 Dec 2025 10:32:53 -0800 Subject: [PATCH 2/4] adjusting tests to work with encoder --- .../MessageEncoderTests.java | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java index 6e83b7f5048c..dde8dc475092 100644 --- a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java +++ b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java @@ -15,18 +15,19 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.Arrays; +import java.util.Objects; import java.util.concurrent.ThreadLocalRandom; import java.util.stream.Stream; +import static com.azure.core.util.FluxUtil.collectBytesInByteBufferStream; +import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.CRC64_LENGTH; +import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_HEADER_LENGTH; +import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_SEGMENT_HEADER_LENGTH; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; public class MessageEncoderTests { - private static final int V1_HEADER_LENGTH = 13; - private static final int V1_SEGMENT_HEADER_LENGTH = 10; - private static final int CRC64_LENGTH = 8; - private static byte[] getRandomData(int size) { byte[] result = new byte[size]; ThreadLocalRandom.current().nextBytes(result); @@ -156,7 +157,7 @@ public void readAll(int size, int segmentSize, StructuredMessageFlags flags) thr StructuredMessageEncoder structuredMessageEncoder = new StructuredMessageEncoder(size, segmentSize, flags); - byte[] actual = structuredMessageEncoder.encode(unencodedBuffer).array(); + byte[] actual = collectBytesInByteBufferStream(structuredMessageEncoder.encode(unencodedBuffer)).block(); byte[] expected = buildStructuredMessage(unencodedBuffer, segmentSize, flags).array(); Assertions.assertArrayEquals(expected, actual); @@ -191,33 +192,39 @@ public void readMultiple(int segmentSize, StructuredMessageFlags flags) throws I byte[] expected = buildStructuredMessage(allWrappedData, segmentSize, flags).array(); ByteArrayOutputStream allActualData = new ByteArrayOutputStream(); - allActualData.write(structuredMessageEncoder.encode(wrappedData1).array()); - allActualData.write(structuredMessageEncoder.encode(wrappedData2).array()); - allActualData.write(structuredMessageEncoder.encode(wrappedData3).array()); + allActualData.write(Objects + .requireNonNull(collectBytesInByteBufferStream(structuredMessageEncoder.encode(wrappedData1)).block())); + allActualData.write(Objects + .requireNonNull(collectBytesInByteBufferStream(structuredMessageEncoder.encode(wrappedData2)).block())); + allActualData.write(Objects + .requireNonNull(collectBytesInByteBufferStream(structuredMessageEncoder.encode(wrappedData3)).block())); Assertions.assertArrayEquals(expected, allActualData.toByteArray()); } @Test - public void emptyBuffer() throws IOException { + public void emptyBuffer() { StructuredMessageEncoder encoder = new StructuredMessageEncoder(10, 5, StructuredMessageFlags.NONE); ByteBuffer emptyBuffer = ByteBuffer.allocate(0); - ByteBuffer result = encoder.encode(emptyBuffer); + ByteBuffer result = ByteBuffer + .wrap(Objects.requireNonNull(collectBytesInByteBufferStream(encoder.encode(emptyBuffer)).block())); assertEquals(0, result.remaining()); } @Test - public void contentAlreadyEncoded() throws IOException { + public void contentAlreadyEncoded() { StructuredMessageEncoder encoder = new StructuredMessageEncoder(4, 2, StructuredMessageFlags.NONE); - encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3, 4 })); - assertThrows(IllegalArgumentException.class, () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 }))); + encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3, 4 })).blockLast(); + assertThrows(IllegalArgumentException.class, + () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })).blockLast()); } @Test - public void bufferLengthExceedsContentLength() throws IOException { + public void bufferLengthExceedsContentLength() { StructuredMessageEncoder encoder = new StructuredMessageEncoder(4, 2, StructuredMessageFlags.NONE); - encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3 })); - assertThrows(IllegalArgumentException.class, () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 }))); + encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3 })).blockLast(); + assertThrows(IllegalArgumentException.class, + () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })).blockLast()); } @Test From 25521844557e414271df6c22566461378f3ccb15 Mon Sep 17 00:00:00 2001 From: Isabelle Date: Wed, 10 Dec 2025 12:54:25 -0800 Subject: [PATCH 3/4] addressing copilot comments --- .../StructuredMessageEncoder.java | 26 +++++++++---------- .../MessageEncoderTests.java | 12 ++++++--- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java index 365773b09631..88bd0b86e8c2 100644 --- a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java +++ b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java @@ -121,27 +121,25 @@ private byte[] generateSegmentHeader(int segmentContentSize) { * * @param unencodedBuffer The buffer to be encoded. * @return A Flux of encoded ByteBuffers. - * @throws IllegalArgumentException If the buffer length exceeds the content length, or the content has already been - * encoded. */ public Flux encode(ByteBuffer unencodedBuffer) { StorageImplUtils.assertNotNull("unencodedBuffer", unencodedBuffer); - if (currentContentOffset == contentLength) { - return Flux - .error(LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded."))); - } + return Flux.defer(() -> { + if (currentContentOffset == contentLength) { + return Flux.error( + LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded."))); + } - if ((unencodedBuffer.remaining() + currentContentOffset) > contentLength) { - return Flux.error( - LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length."))); - } + if ((unencodedBuffer.remaining() + currentContentOffset) > contentLength) { + return Flux.error( + LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length."))); + } - if (!unencodedBuffer.hasRemaining()) { - return Flux.empty(); - } + if (!unencodedBuffer.hasRemaining()) { + return Flux.empty(); + } - return Flux.defer(() -> { List buffers = new ArrayList<>(); // if we are at the beginning of the message, encode message header diff --git a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java index dde8dc475092..aba5ebf2d706 100644 --- a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java +++ b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java @@ -9,6 +9,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import reactor.test.StepVerifier; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -24,6 +25,7 @@ import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_HEADER_LENGTH; import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_SEGMENT_HEADER_LENGTH; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; public class MessageEncoderTests { @@ -215,16 +217,18 @@ public void emptyBuffer() { public void contentAlreadyEncoded() { StructuredMessageEncoder encoder = new StructuredMessageEncoder(4, 2, StructuredMessageFlags.NONE); encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3, 4 })).blockLast(); - assertThrows(IllegalArgumentException.class, - () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })).blockLast()); + StepVerifier.create(encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 }))) + .expectError(IllegalArgumentException.class) + .verify(); } @Test public void bufferLengthExceedsContentLength() { StructuredMessageEncoder encoder = new StructuredMessageEncoder(4, 2, StructuredMessageFlags.NONE); encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2, 3 })).blockLast(); - assertThrows(IllegalArgumentException.class, - () -> encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })).blockLast()); + StepVerifier.create(encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 }))) + .expectError(IllegalArgumentException.class) + .verify(); } @Test From e6d33246d7a78151d01b9dbb99358601cab47601 Mon Sep 17 00:00:00 2001 From: Isabelle Date: Wed, 10 Dec 2025 13:08:26 -0800 Subject: [PATCH 4/4] adding more documentation --- .../structuredmessage/StructuredMessageEncoder.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java index 88bd0b86e8c2..24f80a591948 100644 --- a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java +++ b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java @@ -4,6 +4,7 @@ package com.azure.storage.common.implementation.structuredmessage; import com.azure.core.util.logging.ClientLogger; +import com.azure.storage.common.implementation.BufferStagingArea; import com.azure.storage.common.implementation.StorageCrc64Calculator; import com.azure.storage.common.implementation.StorageImplUtils; import reactor.core.publisher.Flux; @@ -118,9 +119,14 @@ private byte[] generateSegmentHeader(int segmentContentSize) { /** * Encodes the given buffer into a structured message format as a stream of ByteBuffers. + * The encoder maintains mutable state and is designed for single, sequential subscription only. + * Callers should pre-chunk input buffers to appropriate sizes (e.g., using {@link BufferStagingArea}) to + * control memory usage. * * @param unencodedBuffer The buffer to be encoded. * @return A Flux of encoded ByteBuffers. + * @throws IllegalArgumentException If the buffer length exceeds the content length, or the content has already been + * encoded. */ public Flux encode(ByteBuffer unencodedBuffer) { StorageImplUtils.assertNotNull("unencodedBuffer", unencodedBuffer);