Fix smart retry for structured message decoding with relative indexing and byte alignment (#47243)

* Initial plan

* Initial analysis of smart retry issue with structured message decoder

Co-authored-by: gunjansingh-msft <[email protected]>

* Implement smart retry fix for structured message decoding

The fix changes the retry strategy to restart downloads from the beginning
rather than trying to continue mid-stream, which is incompatible with the
structured message format that requires sequential parsing from the start.

Co-authored-by: gunjansingh-msft <[email protected]>

* Add unit tests for DecoderState with byte skipping

Co-authored-by: gunjansingh-msft <[email protected]>

* Update test expectations to match smart retry implementation

All retry requests now start from offset 0 for structured message validation
since the decoder must parse complete messages from the beginning.

Co-authored-by: gunjansingh-msft <[email protected]>

* Fix smart retry to properly track consumed encoded bytes

The original implementation tracked received buffer sizes before decoding,
which didn't account for pending bytes from previous buffers. This caused
retry offsets to be incorrect. Now tracking only newly consumed bytes
after accounting for pending data.

Co-authored-by: gunjansingh-msft <[email protected]>

* Revert test expectations and unused constants

Since smart retry now correctly resumes from interrupted offset (not from
beginning), reverted test expectations to validate progressive offsets.
Also removed unused DECODED_BYTES_TO_SKIP constant and unit tests that
were specific to the "restart from beginning" approach.

Co-authored-by: gunjansingh-msft <[email protected]>

* Fix byte tracking to count received bytes not consumed bytes

The issue was that totalEncodedBytesProcessed was tracking consumed bytes
(what the decoder processed), but it should track received bytes (what came
from the network). This ensures retry requests fetch new data after what
was already received, while the pending buffer holds unconsumed bytes.

On retry: request starts after all received bytes, decoder continues from
its internal state with pending bytes, new bytes appended after pending.

Co-authored-by: gunjansingh-msft <[email protected]>

* Fix retry offset calculation to use getRetryOffset() method

The critical bug was using getTotalEncodedBytesProcessed() for retry offset,
which includes pending bytes. This caused gaps in the byte stream when
combining pending + new data. The existing getRetryOffset() method correctly
subtracts pending bytes to get the actual file position where new data
should start, ensuring continuous byte streams across retries.

Co-authored-by: gunjansingh-msft <[email protected]>

* Track segment boundaries for smart retry resumption

Implemented segment boundary tracking in StructuredMessageDecoder:
- Added lastCompleteSegmentStart to track where complete segments end
- Changed messageOffset and segment fields to long for proper offset tracking
- Fixed segment size reading to handle unsigned long values correctly
- Updated DecoderState.getRetryOffset() to use segment boundaries
- Ensures retries resume from complete segment boundaries, not mid-segment

This matches the decryptor pattern where retries align with block boundaries.

Co-authored-by: gunjansingh-msft <[email protected]>

* Add debug logging for segment boundary tracking and retry offsets

Added verbose logging to track:
- Segment completion with byte offsets in StructuredMessageDecoder
- Retry offset calculation in DecoderState
- Retry range in BlobAsyncClientBase for structured messages

This helps diagnose smart retry behavior and verify segment boundary alignment.

Co-authored-by: gunjansingh-msft <[email protected]>

* Reset decoder and clear pending buffer on retry

Critical fix for smart retry:
- Added resetToLastCompleteSegment() to decoder to rewind messageOffset
- Call reset in getRetryOffset() to sync decoder state with retry offset
- Clear pending buffer since bytes after segment boundary will be re-fetched
- Ensures decoder's messageOffset matches the file position being requested

This prevents decoder from reading at wrong offsets after retry.

Co-authored-by: gunjansingh-msft <[email protected]>

* Add comprehensive diagnostic logging for smart retry debugging

Added detailed logging to track decoder state and buffer positions:
- readAndValidateSegmentLength() with hex dump of buffer contents
- Enhanced resetToLastCompleteSegment() logging with all state details
- Segment completion logging with offset and length details
- Buffer reception logging in decodeStream with pending/new byte counts
- Retry offset calculation logging with before/after decoder state

This will help identify where decoder state and retry offset become misaligned.

Co-authored-by: gunjansingh-msft <[email protected]>

* Fix byte alignment and endianness in decoder buffer handling

- combineWithPending: returns buffer with position=0 and LITTLE_ENDIAN
- updatePendingBuffer: allocates buffer with LITTLE_ENDIAN order
- decodeStream: proper position tracking for consumed bytes
- Added debug logging for decode iteration

Co-authored-by: gunjansingh-msft <[email protected]>

* Fix decodeStream with relative indexing for correct segment alignment

Implemented the correct algorithm for smart retry:
- Capture absoluteStartOfCombined before adding new bytes
- Calculate relativeIndex = decoderOffset - absoluteStartOfCombined
- Add peekNextSegmentLength to peek segment length without consuming
- Compute encodedSegmentSize = header + payload + CRC
- Slice complete segments only before decoding
- Update totalEncodedBytesProcessed after each segment is decoded

This ensures proper byte alignment and segment boundary handling.

Co-authored-by: gunjansingh-msft <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: gunjansingh-msft <[email protected]>

Author: Copilot

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/specialized/BlobAsyncClientBase.java

@@ -1407,13 +1407,16 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                             if (decoderStateObj instanceof StorageContentValidationDecoderPolicy.DecoderState) {
                                 DecoderState decoderState = (DecoderState) decoderStateObj;
 
-                                // Use totalEncodedBytesProcessed to request NEW bytes from the server
-                                // The pending buffer already contains bytes we've received, so we request
-                                // starting from the next byte after what we've already received
-                                long encodedOffset = decoderState.getTotalEncodedBytesProcessed();
+                                // Use getRetryOffset() to get the correct offset for retry
+                                // This accounts for pending bytes that have been received but not yet consumed
+                                long encodedOffset = decoderState.getRetryOffset();
                                 long remainingCount = finalCount - encodedOffset;
                                 retryRange = new BlobRange(initialOffset + encodedOffset, remainingCount);
 
+                                LOGGER.info(
+                                    "Structured message smart retry: resuming from offset {} (initial={}, encoded={})",
+                                    initialOffset + encodedOffset, initialOffset, encodedOffset);
+
                                 // Preserve the decoder state for the retry
                                 retryContext = retryContext
                                     .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageDecoder.java

@@ -27,15 +27,19 @@ public class StructuredMessageDecoder {
     private int numSegments;
     private final long expectedContentLength;
 
-    private int messageOffset = 0;
+    private long messageOffset = 0;
     private int currentSegmentNumber = 0;
-    private int currentSegmentContentLength = 0;
-    private int currentSegmentContentOffset = 0;
+    private long currentSegmentContentLength = 0;
+    private long currentSegmentContentOffset = 0;
 
     private long messageCrc64 = 0;
     private long segmentCrc64 = 0;
     private final Map<Integer, Long> segmentCrcs = new HashMap<>();
 
+    // Track the last complete segment boundary for smart retry
+    private long lastCompleteSegmentStart = 0;
+    private long currentSegmentStart = 0;
+
     /**
      * Constructs a new StructuredMessageDecoder.
      *
@@ -45,6 +49,50 @@ public StructuredMessageDecoder(long expectedContentLength) {
         this.expectedContentLength = expectedContentLength;
     }
 
+    /**
+     * Gets the byte offset where the last complete segment ended.
+     * This is used for smart retry to resume from a segment boundary.
+     *
+     * @return The byte offset of the last complete segment boundary.
+     */
+    public long getLastCompleteSegmentStart() {
+        return lastCompleteSegmentStart;
+    }
+
+    /**
+     * Gets the current message offset (total bytes consumed from the structured message).
+     *
+     * @return The current message offset.
+     */
+    public long getMessageOffset() {
+        return messageOffset;
+    }
+
+    /**
+     * Resets the decoder position to the last complete segment boundary.
+     * This is used during smart retry to ensure the decoder is in sync with
+     * the data being provided from the retry offset.
+     */
+    public void resetToLastCompleteSegment() {
+        if (messageOffset != lastCompleteSegmentStart) {
+            LOGGER.atInfo()
+                .addKeyValue("fromOffset", messageOffset)
+                .addKeyValue("toOffset", lastCompleteSegmentStart)
+                .addKeyValue("currentSegmentNum", currentSegmentNumber)
+                .addKeyValue("currentSegmentContentOffset", currentSegmentContentOffset)
+                .addKeyValue("currentSegmentContentLength", currentSegmentContentLength)
+                .log("Resetting decoder to last complete segment boundary");
+            messageOffset = lastCompleteSegmentStart;
+            // Reset current segment state - next decode will read the segment header
+            currentSegmentContentOffset = 0;
+            currentSegmentContentLength = 0;
+        } else {
+            LOGGER.atVerbose()
+                .addKeyValue("offset", messageOffset)
+                .log("Decoder already at last complete segment boundary, no reset needed");
+        }
+    }
+
     /**
      * Reads the message header from the given buffer.
      *
@@ -79,6 +127,91 @@ private void readMessageHeader(ByteBuffer buffer) {
         messageOffset += V1_HEADER_LENGTH;
     }
 
+    /**
+     * Converts a ByteBuffer range to hex string for diagnostic purposes.
+     */
+    private static String toHex(ByteBuffer buf, int len) {
+        int pos = buf.position();
+        int peek = Math.min(len, buf.remaining());
+        byte[] out = new byte[peek];
+        buf.get(out, 0, peek);
+        buf.position(pos);
+        StringBuilder sb = new StringBuilder();
+        for (int i = 0; i < out.length; i++) {
+            sb.append(String.format("%02X", out[i]));
+            if (i < out.length - 1)
+                sb.append(' ');
+        }
+        return sb.toString();
+    }
+
+    /**
+     * Peeks the next segment length without consuming from the buffer.
+     * Used by the policy to calculate encoded segment size before slicing.
+     *
+     * @param buffer The buffer to peek from.
+     * @param relativeIndex The position in the buffer to start reading from.
+     * @return The segment content length, or -1 if not enough bytes.
+     */
+    public long peekNextSegmentLength(ByteBuffer buffer, int relativeIndex) {
+        // Need at least V1_SEGMENT_HEADER_LENGTH bytes to read segment number (2) + segment size (8)
+        if (relativeIndex + V1_SEGMENT_HEADER_LENGTH > buffer.limit()) {
+            return -1;
+        }
+        // Segment size is at offset 2 (after segment number which is 2 bytes)
+        return buffer.getLong(relativeIndex + 2);
+    }
+
+    /**
+     * Gets the flags for the current message (needed to determine if CRC is present).
+     *
+     * @return The message flags, or null if header not yet read.
+     */
+    public StructuredMessageFlags getFlags() {
+        return flags;
+    }
+
+    /**
+     * Reads and validates segment length with diagnostic logging.
+     */
+    private long readAndValidateSegmentLength(ByteBuffer buffer, long remaining) {
+        final int SEGMENT_SIZE_BYTES = 8; // segment size is 8 bytes (long)
+        if (buffer.remaining() < SEGMENT_SIZE_BYTES) {
+            LOGGER.error("Not enough bytes to read segment size. pos={}, remaining={}", buffer.position(),
+                buffer.remaining());
+            throw new IllegalStateException("Not enough bytes to read segment size");
+        }
+
+        // Diagnostic: dump first 16 bytes at this position so we can see what's being read
+        LOGGER.atInfo()
+            .addKeyValue("decoderOffset", messageOffset)
+            .addKeyValue("bufferPos", buffer.position())
+            .addKeyValue("bufferRemaining", buffer.remaining())
+            .addKeyValue("peek16", toHex(buffer, 16))
+            .addKeyValue("lastCompleteSegment", lastCompleteSegmentStart)
+            .log("Decoder about to read segment length");
+
+        long segmentLength = buffer.getLong();
+
+        if (segmentLength < 0 || segmentLength > remaining) {
+            // Peek next bytes for extra detail
+            String peekNext = toHex(buffer, 16);
+            LOGGER.error(
+                "Invalid segment length read: segmentLength={}, remaining={}, decoderOffset={}, "
+                    + "lastCompleteSegment={}, bufferPos={}, peek-next-bytes={}",
+                segmentLength, remaining, messageOffset, lastCompleteSegmentStart, buffer.position(), peekNext);
+            throw new IllegalArgumentException("Invalid segment size detected: " + segmentLength + " (remaining="
+                + remaining + ", decoderOffset=" + messageOffset + ")");
+        }
+
+        LOGGER.atVerbose()
+            .addKeyValue("segmentLength", segmentLength)
+            .addKeyValue("decoderOffset", messageOffset)
+            .log("Valid segment length read");
+
+        return segmentLength;
+    }
+
     /**
      * Reads the segment header from the given buffer.
      *
@@ -90,13 +223,13 @@ private void readSegmentHeader(ByteBuffer buffer) {
             throw LOGGER.logExceptionAsError(new IllegalArgumentException("Segment header is incomplete."));
         }
 
+        // Mark the start of this segment (before reading the header)
+        currentSegmentStart = messageOffset;
+
         int segmentNum = Short.toUnsignedInt(buffer.getShort());
-        int segmentSize = (int) buffer.getLong();
 
-        if (segmentSize < 0 || segmentSize > buffer.remaining()) {
-            throw LOGGER
-                .logExceptionAsError(new IllegalArgumentException("Invalid segment size detected: " + segmentSize));
-        }
+        // Read segment size with validation and diagnostics
+        long segmentSize = readAndValidateSegmentLength(buffer, buffer.remaining());
 
         if (segmentNum != currentSegmentNumber + 1) {
             throw LOGGER.logExceptionAsError(new IllegalArgumentException("Unexpected segment number."));
@@ -126,8 +259,8 @@ private void readSegmentHeader(ByteBuffer buffer) {
      * @throws IllegalArgumentException if there is a segment size mismatch.
      */
     private void readSegmentContent(ByteBuffer buffer, ByteArrayOutputStream output, int size) {
-        int toRead = Math.min(buffer.remaining(), currentSegmentContentLength - currentSegmentContentOffset);
-        toRead = Math.min(toRead, size);
+        long remaining = currentSegmentContentLength - currentSegmentContentOffset;
+        int toRead = (int) Math.min(buffer.remaining(), Math.min(remaining, size));
 
         if (toRead == 0) {
             return;
@@ -182,10 +315,17 @@ private void readSegmentFooter(ByteBuffer buffer) {
             messageOffset += CRC64_LENGTH;
         }
 
+        // Mark that this segment is complete - update the last complete segment boundary
+        // This is the position where we can safely resume if a retry occurs
+        lastCompleteSegmentStart = messageOffset;
+        LOGGER.atInfo()
+            .addKeyValue("segmentNum", currentSegmentNumber)
+            .addKeyValue("offset", lastCompleteSegmentStart)
+            .addKeyValue("segmentLength", currentSegmentContentLength)
+            .log("Segment complete at byte offset");
+
         if (currentSegmentNumber == numSegments) {
             readMessageFooter(buffer);
-        } else {
-            readSegmentHeader(buffer);
         }
     }