Fix smart retry for structured message decoding byte tracking

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

@@ -1395,31 +1395,28 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                         Context retryContext = firstRangeContext;
                         BlobRange retryRange;
 
-                        // If structured message decoding is enabled, we need to calculate the retry offset
-                        // based on the encoded bytes processed, not the decoded bytes
+                        // If structured message decoding is enabled, we need to restart from the beginning
+                        // because the decoder must parse the complete structured message from the start
                         if (contentValidationOptions != null
                             && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                            // Get the decoder state to determine how many encoded bytes were processed
+                            // Get the decoder state to determine how many decoded bytes were already emitted
                             Object decoderStateObj
                                 = firstRangeContext.getData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY)
                                     .orElse(null);
 
+                            // For structured message validation, we must restart from the beginning
+                            // because the message has headers and sequential segment numbers that must
+                            // be parsed in order. We cannot resume parsing mid-stream.
+                            retryRange = new BlobRange(initialOffset, finalCount);
+
+                            // DO NOT preserve decoder state - create a fresh decoder for the retry
+                            // The policy will track how many decoded bytes to skip
                             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();
-                                long remainingCount = finalCount - encodedOffset;
-                                retryRange = new BlobRange(initialOffset + encodedOffset, remainingCount);
-
-                                // Preserve the decoder state for the retry
-                                retryContext = retryContext
-                                    .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
-                            } else {
-                                // No decoder state yet, use the normal retry logic
-                                retryRange = new BlobRange(initialOffset + offset, newCount);
+                                // Add the current decoded offset so the policy knows how many bytes to skip
+                                retryContext = retryContext.addData(
+                                    Constants.STRUCTURED_MESSAGE_DECODED_BYTES_TO_SKIP_CONTEXT_KEY,
+                                    decoderState.getTotalBytesDecoded());
                             }
                         } else {
                             // For non-structured downloads, use smart retry from the interrupted offset

sdk/storage/azure-storage-blob/src/test/java/com/azure/storage/blob/specialized/BlobMessageDecoderDownloadTests.java

@@ -272,19 +272,14 @@ public void downloadStreamWithResponseContentValidationSmartRetry() throws IOExc
         List<String> rangeHeaders = mockPolicy.getRangeHeaders();
         assertTrue(rangeHeaders.size() > 0, "Expected range headers for retries");
 
-        // With structured message validation and smart retry, retries should resume from the encoded
-        // offset where the interruption occurred. The first request starts at 0, and subsequent
-        // retry requests should start from progressively higher offsets.
-        assertTrue(rangeHeaders.get(0).startsWith("bytes=0-"), "First request should start from offset 0");
-
-        // Subsequent requests should start from higher offsets (smart retry resuming from where it left off)
-        for (int i = 1; i < rangeHeaders.size(); i++) {
+        // With structured message validation and smart retry, all requests (initial and retries)
+        // must start from offset 0 because the structured message format requires sequential
+        // parsing from the beginning. The decoder cannot resume mid-stream.
+        // The policy skips already-emitted decoded bytes to avoid duplication.
+        for (int i = 0; i < rangeHeaders.size(); i++) {
             String rangeHeader = rangeHeaders.get(i);
-            // Each retry should start from a higher offset than the previous
-            // Note: We can't assert exact offset values as they depend on how much data was received
-            // before the interruption, but we can verify it's a valid range header
-            assertTrue(rangeHeader.startsWith("bytes="),
-                "Retry request " + i + " should have a range header: " + rangeHeader);
+            assertTrue(rangeHeader.startsWith("bytes=0-"),
+                "Request " + i + " should start from offset 0 for structured message validation: " + rangeHeader);
         }
     }
 
@@ -334,18 +329,18 @@ public void downloadStreamWithResponseContentValidationSmartRetryMultipleSegment
         assertTrue(rangeHeaders.size() >= 4,
             "Expected at least 4 range headers for retries, got: " + rangeHeaders.size());
 
-        // With smart retry, each request should have a valid range header
+        // With smart retry and structured message validation, all requests must start from offset 0
         for (int i = 0; i < rangeHeaders.size(); i++) {
             String rangeHeader = rangeHeaders.get(i);
-            assertTrue(rangeHeader.startsWith("bytes="),
-                "Request " + i + " should have a valid range header, but was: " + rangeHeader);
+            assertTrue(rangeHeader.startsWith("bytes=0-"), "Request " + i
+                + " should start from offset 0 for structured message validation, but was: " + rangeHeader);
         }
     }
 
     @Test
     public void downloadStreamWithResponseContentValidationSmartRetryLargeBlob() throws IOException {
-        // Test smart retry with a larger blob to ensure retries resume from the
-        // interrupted offset and successfully validate all data
+        // Test smart retry with a larger blob to ensure retries restart from the beginning
+        // and successfully validate all data after skipping already-emitted decoded bytes
 
         byte[] randomData = getRandomByteArray(5 * Constants.KB);
         StructuredMessageEncoder encoder
@@ -383,12 +378,12 @@ public void downloadStreamWithResponseContentValidationSmartRetryLargeBlob() thr
         // Verify that retries occurred
         assertEquals(0, mockPolicy.getTriesRemaining());
 
-        // Verify that smart retry is working with valid range headers
+        // Verify that all requests start from offset 0 for structured message validation
         List<String> rangeHeaders = mockPolicy.getRangeHeaders();
         for (int i = 0; i < rangeHeaders.size(); i++) {
             String rangeHeader = rangeHeaders.get(i);
-            assertTrue(rangeHeader.startsWith("bytes="),
-                "Request " + i + " should have a valid range header, but was: " + rangeHeader);
+            assertTrue(rangeHeader.startsWith("bytes=0-"), "Request " + i
+                + " should start from offset 0 for structured message validation, but was: " + rangeHeader);
         }
     }
 }

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

@@ -111,6 +111,14 @@ public final class Constants {
     public static final String STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY
         = "azure-storage-structured-message-decoder-state";
 
+    /**
+     * Context key used to pass the number of decoded bytes to skip on retry.
+     * When a retry occurs with structured message validation, we restart from the beginning
+     * and skip bytes that were already emitted to avoid duplication.
+     */
+    public static final String STRUCTURED_MESSAGE_DECODED_BYTES_TO_SKIP_CONTEXT_KEY
+        = "azure-storage-structured-message-decoded-bytes-to-skip";
+
     private Constants() {
     }
 

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

@@ -60,8 +60,15 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
             Long contentLength = getContentLength(httpResponse.getHeaders());
 
             if (contentLength != null && contentLength > 0 && validationOptions != null) {
-                // Get or create decoder with state tracking
-                DecoderState decoderState = getOrCreateDecoderState(context, contentLength);
+                // Check if this is a retry - if so, get the number of decoded bytes to skip
+                long bytesToSkip = context.getData(Constants.STRUCTURED_MESSAGE_DECODED_BYTES_TO_SKIP_CONTEXT_KEY)
+                    .filter(value -> value instanceof Long)
+                    .map(value -> (Long) value)
+                    .orElse(0L);
+
+                // Always create a fresh decoder for each request
+                // This is necessary because structured messages must be parsed from the beginning
+                DecoderState decoderState = new DecoderState(contentLength, bytesToSkip);
 
                 // Decode using the stateful decoder
                 Flux<ByteBuffer> decodedStream = decodeStream(httpResponse.getBody(), decoderState);
@@ -98,14 +105,10 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                 int availableSize = dataToProcess.remaining();
                 ByteBuffer duplicateForDecode = dataToProcess.duplicate();
                 int initialPosition = duplicateForDecode.position();
-                
+
                 // Decode - this advances duplicateForDecode's position
                 ByteBuffer decodedData = state.decoder.decode(duplicateForDecode, availableSize);
 
-                // Track decoded bytes
-                int decodedBytes = decodedData.remaining();
-                state.totalBytesDecoded.addAndGet(decodedBytes);
-
                 // Calculate how much of the input buffer was consumed by checking the duplicate's position
                 int bytesConsumed = duplicateForDecode.position() - initialPosition;
                 int bytesRemaining = availableSize - bytesConsumed;
@@ -121,9 +124,33 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                     state.pendingBuffer = null;
                 }
 
-                // Return decoded data if any
+                // Handle skipping bytes for retries and tracking decoded bytes
+                int decodedBytes = decodedData.remaining();
                 if (decodedBytes > 0) {
-                    return Flux.just(decodedData);
+                    // Track total decoded bytes
+                    long totalDecoded = state.totalBytesDecoded.addAndGet(decodedBytes);
+
+                    // If we need to skip bytes (retry scenario), adjust the buffer
+                    if (state.bytesToSkip > 0) {
+                        long currentPosition = totalDecoded - decodedBytes; // Where we were before adding these bytes
+
+                        if (currentPosition + decodedBytes <= state.bytesToSkip) {
+                            // All these bytes should be skipped
+                            return Flux.empty();
+                        } else if (currentPosition < state.bytesToSkip) {
+                            // Some bytes should be skipped
+                            int skipAmount = (int) (state.bytesToSkip - currentPosition);
+                            decodedData.position(decodedData.position() + skipAmount);
+                        }
+                        // else: no bytes need to be skipped, emit all
+                    }
+
+                    // Return decoded data if any remains after skipping
+                    if (decodedData.hasRemaining()) {
+                        return Flux.just(decodedData);
+                    } else {
+                        return Flux.empty();
+                    }
                 } else {
                     return Flux.empty();
                 }
@@ -199,20 +226,6 @@ private Long getContentLength(HttpHeaders headers) {
         return null;
     }
 
-    /**
-     * Gets or creates a decoder state from context.
-     *
-     * @param context The pipeline call context.
-     * @param contentLength The content length.
-     * @return The decoder state.
-     */
-    private DecoderState getOrCreateDecoderState(HttpPipelineCallContext context, long contentLength) {
-        return context.getData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY)
-            .filter(value -> value instanceof DecoderState)
-            .map(value -> (DecoderState) value)
-            .orElseGet(() -> new DecoderState(contentLength));
-    }
-
     /**
      * Checks if the response is a download response.
      *
@@ -233,18 +246,21 @@ public static class DecoderState {
         private final long expectedContentLength;
         private final AtomicLong totalBytesDecoded;
         private final AtomicLong totalEncodedBytesProcessed;
+        private final long bytesToSkip;
         private ByteBuffer pendingBuffer;
 
         /**
          * Creates a new decoder state.
          *
          * @param expectedContentLength The expected length of the encoded content.
+         * @param bytesToSkip The number of decoded bytes to skip (for retry scenarios).
          */
-        public DecoderState(long expectedContentLength) {
+        public DecoderState(long expectedContentLength, long bytesToSkip) {
             this.expectedContentLength = expectedContentLength;
             this.decoder = new StructuredMessageDecoder(expectedContentLength);
             this.totalBytesDecoded = new AtomicLong(0);
             this.totalEncodedBytesProcessed = new AtomicLong(0);
+            this.bytesToSkip = bytesToSkip;
             this.pendingBuffer = null;
         }
 

sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicyTest.java

@@ -3,9 +3,14 @@
 
 package com.azure.storage.common.policy;
 
+import com.azure.storage.common.policy.StorageContentValidationDecoderPolicy.DecoderState;
 import org.junit.jupiter.api.Test;
 
+import java.nio.ByteBuffer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
  * Unit tests for {@link StorageContentValidationDecoderPolicy}.
@@ -21,4 +26,29 @@ public void policyCanBeInstantiated() {
         StorageContentValidationDecoderPolicy policy = new StorageContentValidationDecoderPolicy();
         assertNotNull(policy);
     }
+
+    @Test
+    public void decoderStateTracksDecodedBytes() {
+        // Create a decoder state with no bytes to skip
+        DecoderState state = new DecoderState(1024, 0);
+
+        assertNotNull(state);
+        assertEquals(0, state.getTotalBytesDecoded());
+        assertEquals(0, state.getTotalEncodedBytesProcessed());
+    }
+
+    @Test
+    public void decoderStateWithBytesToSkip() {
+        // Create a decoder state with bytes to skip (retry scenario)
+        long bytesToSkip = 512;
+        DecoderState state = new DecoderState(1024, bytesToSkip);
+
+        assertNotNull(state);
+        assertEquals(0, state.getTotalBytesDecoded());
+        assertEquals(0, state.getTotalEncodedBytesProcessed());
+
+        // Verify the bytesToSkip field is set correctly
+        // Note: We can't directly access bytesToSkip as it's private, 
+        // but it will be used internally during decoding
+    }
 }