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

Author: Copilot

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

@@ -1395,28 +1395,31 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                         Context retryContext = firstRangeContext;
                         BlobRange retryRange;
 
-                        // 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 structured message decoding is enabled, we need to calculate the retry offset
+                        // based on the encoded bytes processed, not the decoded bytes
                         if (contentValidationOptions != null
                             && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                            // Get the decoder state to determine how many decoded bytes were already emitted
+                            // Get the decoder state to determine how many encoded bytes were processed
                             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;
-                                // 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());
+
+                                // 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);
                             }
                         } else {
                             // For non-structured downloads, use smart retry from the interrupted offset

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

@@ -60,15 +60,8 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
             Long contentLength = getContentLength(httpResponse.getHeaders());
 
             if (contentLength != null && contentLength > 0 && validationOptions != null) {
-                // 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);
+                // Get or create decoder with state tracking
+                DecoderState decoderState = getOrCreateDecoderState(context, contentLength);
 
                 // Decode using the stateful decoder
                 Flux<ByteBuffer> decodedStream = decodeStream(httpResponse.getBody(), decoderState);
@@ -92,13 +85,12 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
      */
     private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState state) {
         return encodedFlux.concatMap(encodedBuffer -> {
+            // Track how many bytes were pending before we process
+            int previousPendingBytes = (state.pendingBuffer != null) ? state.pendingBuffer.remaining() : 0;
+
             // Combine with pending data if any
             ByteBuffer dataToProcess = state.combineWithPending(encodedBuffer);
 
-            // Track encoded bytes
-            int encodedBytesInBuffer = encodedBuffer.remaining();
-            state.totalEncodedBytesProcessed.addAndGet(encodedBytesInBuffer);
-
             try {
                 // Try to decode what we have - decoder handles partial data
                 // Create duplicate for decoder - it will advance the duplicate's position as it reads
@@ -113,6 +105,14 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                 int bytesConsumed = duplicateForDecode.position() - initialPosition;
                 int bytesRemaining = availableSize - bytesConsumed;
 
+                // Track the newly consumed encoded bytes (excluding previously pending bytes)
+                // The consumed bytes include both old pending bytes and new bytes from this buffer
+                // We only want to add the NEW bytes that were consumed
+                int newBytesConsumed = bytesConsumed - previousPendingBytes;
+                if (newBytesConsumed > 0) {
+                    state.totalEncodedBytesProcessed.addAndGet(newBytesConsumed);
+                }
+
                 // Save only unconsumed portion to pending
                 if (bytesRemaining > 0) {
                     // Position the original buffer to skip consumed bytes, then slice to get unconsumed
@@ -124,33 +124,13 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                     state.pendingBuffer = null;
                 }
 
-                // Handle skipping bytes for retries and tracking decoded bytes
+                // Track decoded bytes
                 int decodedBytes = decodedData.remaining();
+                state.totalBytesDecoded.addAndGet(decodedBytes);
+
+                // Return decoded data if any
                 if (decodedBytes > 0) {
-                    // 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();
-                    }
+                    return Flux.just(decodedData);
                 } else {
                     return Flux.empty();
                 }
@@ -226,6 +206,20 @@ 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.
      *
@@ -246,21 +240,18 @@ 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, long bytesToSkip) {
+        public DecoderState(long expectedContentLength) {
             this.expectedContentLength = expectedContentLength;
             this.decoder = new StructuredMessageDecoder(expectedContentLength);
             this.totalBytesDecoded = new AtomicLong(0);
             this.totalEncodedBytesProcessed = new AtomicLong(0);
-            this.bytesToSkip = bytesToSkip;
             this.pendingBuffer = null;
         }