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

Author: Copilot

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-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);
@@ -102,10 +109,6 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                 // 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;
         }