Skip to content

Storage Content Validation - Encoder Performance Improvements #47531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ibrandes merged 4 commits into from
Dec 14, 2025
Merged

Storage Content Validation - Encoder Performance Improvements #47531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
19cc55c
adjusting encoder logic
ibrandes Dec 10, 2025
753f5b5
adjusting tests to work with encoder
ibrandes Dec 10, 2025
2552184
addressing copilot comments
ibrandes Dec 10, 2025
e6d3324
adding more documentation
ibrandes Dec 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 65 additions & 79 deletions ...a/com/azure/storage/common/implementation/structuredmessage/StructuredMessageEncoder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,16 @@
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;

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;
Expand All @@ -35,7 +37,6 @@ public class StructuredMessageEncoder {
private int currentContentOffset;
private int currentSegmentNumber;
private int currentSegmentOffset;
private int currentMessageLength;
private long messageCRC64;
private final Map<Integer, Long> segmentCRC64s;

Expand Down Expand Up @@ -66,7 +67,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);
Expand Down Expand Up @@ -109,115 +109,102 @@ 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.
* 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 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<ByteBuffer> encode(ByteBuffer unencodedBuffer) {
StorageImplUtils.assertNotNull("unencodedBuffer", unencodedBuffer);

if (currentContentOffset == contentLength) {
throw LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded."));
}

if ((unencodedBuffer.remaining() + currentContentOffset) > contentLength) {
throw LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length."));
}

if (!unencodedBuffer.hasRemaining()) {
return ByteBuffer.allocate(0);
}

ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();

// 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);
return Flux.defer(() -> {
if (currentContentOffset == contentLength) {
return Flux.error(
LOGGER.logExceptionAsError(new IllegalArgumentException("Content has already been encoded.")));
}

encodeSegmentContent(unencodedBuffer, byteArrayOutputStream);

// if we are at the end of a segment's content, encode segment footer
if (currentSegmentOffset == getSegmentContentLength()) {
encodeSegmentFooter(byteArrayOutputStream);
if ((unencodedBuffer.remaining() + currentContentOffset) > contentLength) {
return Flux.error(
LOGGER.logExceptionAsError(new IllegalArgumentException("Buffer length exceeds content length.")));
}
}

// if all content has been encoded, encode message footer
if (currentContentOffset == contentLength) {
encodeMessageFooter(byteArrayOutputStream);
}
if (!unencodedBuffer.hasRemaining()) {
return Flux.empty();
}

return ByteBuffer.wrap(byteArrayOutputStream.toByteArray());
}
List<ByteBuffer> buffers = new ArrayList<>();

private void encodeMessageHeader(ByteArrayOutputStream output) {
byte[] metadata = generateMessageHeader();
output.write(metadata, 0, metadata.length);
// if we are at the beginning of the message, encode message header
if (currentContentOffset == 0) {
buffers.add(ByteBuffer.wrap(generateMessageHeader()));
}

currentMessageLength += metadata.length;
}
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;
}
}

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);

Expand All @@ -230,8 +217,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() {
Expand All @@ -255,7 +241,7 @@ private void incrementCurrentSegment() {
*
* @return The length of the message.
*/
public int getMessageLength() {
public long getEncodedMessageLength() {
Comment thread
ibrandes marked this conversation as resolved.
return messageLength;
}
}
43 changes: 27 additions & 16 deletions ...t/java/com/azure/storage/common/implementation/structuredmessage/MessageEncoderTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,27 @@
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;
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.assertNotNull;
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);
Expand Down Expand Up @@ -156,7 +159,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);
Expand Down Expand Up @@ -191,33 +194,41 @@ 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()));
Comment thread
ibrandes marked this conversation as resolved.
allActualData.write(Objects
.requireNonNull(collectBytesInByteBufferStream(structuredMessageEncoder.encode(wrappedData2)).block()));
allActualData.write(Objects
.requireNonNull(collectBytesInByteBufferStream(structuredMessageEncoder.encode(wrappedData3)).block()));
Comment thread
ibrandes marked this conversation as resolved.

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();
StepVerifier.create(encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })))
.expectError(IllegalArgumentException.class)
.verify();
}

@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();
StepVerifier.create(encoder.encode(ByteBuffer.wrap(new byte[] { 1, 2 })))
.expectError(IllegalArgumentException.class)
.verify();
}

@Test
Expand Down
Loading