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]>

Author: Copilot

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,25 @@ 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;
+    }
+
     /**
      * Reads the message header from the given buffer.
      *
@@ -90,10 +113,16 @@ 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()) {
+        // Read segment size as long (8 bytes)
+        long segmentSize = buffer.getLong();
+
+        // Validate segment size
+        if (segmentSize < 0L || segmentSize > buffer.remaining()) {
             throw LOGGER
                 .logExceptionAsError(new IllegalArgumentException("Invalid segment size detected: " + segmentSize));
         }
@@ -126,8 +155,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 +211,12 @@ 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;
+
         if (currentSegmentNumber == numSegments) {
             readMessageFooter(buffer);
-        } else {
-            readSegmentHeader(buffer);
         }
     }
 

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java

@@ -297,15 +297,15 @@ public long getTotalEncodedBytesProcessed() {
 
         /**
          * Gets the offset to use for retry requests.
-         * This is the total encoded bytes processed minus any bytes in the pending buffer,
-         * since pending bytes have already been counted but haven't been successfully processed yet.
+         * This uses the decoder's last complete segment boundary to ensure retries
+         * resume from a valid segment boundary, not mid-segment.
          *
-         * @return The offset for retry requests.
+         * @return The offset for retry requests (last complete segment boundary).
          */
         public long getRetryOffset() {
-            long processed = totalEncodedBytesProcessed.get();
-            int pending = (pendingBuffer != null) ? pendingBuffer.remaining() : 0;
-            return processed - pending;
+            // Use the decoder's last complete segment start as the retry offset
+            // This ensures we resume from a segment boundary, not mid-segment
+            return decoder.getLastCompleteSegmentStart();
         }
 
         /**