smart retry changes

gunjansingh-msft · gunjansingh-msft · commit a9abd2ee249e · 2025-10-29T19:47:57.000+05:30
diff --git a/sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/specialized/BlobAsyncClientBase.java b/sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/specialized/BlobAsyncClientBase.java
@@ -1340,8 +1340,26 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                     }
 
                     try {
+                        // For retry context, preserve decoder state if structured message validation is enabled
+                        Context retryContext = firstRangeContext;
+
+                        // If structured message decoding is enabled, we need to include the decoder state
+                        // so the retry can continue from where we left off
+                        if (contentValidationOptions != null
+                            && contentValidationOptions.isStructuredMessageValidationEnabled()) {
+                            // The decoder state will be set by the policy during processing
+                            // We preserve it in the context for the retry request
+                            Object decoderState
+                                = firstRangeContext.getData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY)
+                                    .orElse(null);
+                            if (decoderState != null) {
+                                retryContext = retryContext
+                                    .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
+                            }
+                        }
+
                         return downloadRange(new BlobRange(initialOffset + offset, newCount), finalRequestConditions,
-                            eTag, finalGetMD5, firstRangeContext);
+                            eTag, finalGetMD5, retryContext);
                     } catch (Exception e) {
                         return Mono.error(e);
                     }
diff --git a/sdk/storage/azure-storage-blob/src/test/java/com/azure/storage/blob/BlobMessageDecoderDownloadTests.java b/sdk/storage/azure-storage-blob/src/test/java/com/azure/storage/blob/BlobMessageDecoderDownloadTests.java
@@ -6,6 +6,8 @@
 import com.azure.core.test.utils.TestUtils;
 import com.azure.core.util.FluxUtil;
 import com.azure.storage.blob.models.BlobRange;
+import com.azure.storage.blob.models.BlobRequestConditions;
+import com.azure.storage.blob.models.DownloadRetryOptions;
 import com.azure.storage.common.DownloadContentValidationOptions;
 import com.azure.storage.common.implementation.Constants;
 import com.azure.storage.common.implementation.structuredmessage.StructuredMessageEncoder;
@@ -18,6 +20,7 @@
 import java.io.IOException;
 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;
 
diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/Constants.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/Constants.java
@@ -105,6 +105,12 @@ public final class Constants {
     public static final String STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY
         = "azure-storage-structured-message-validation-options";
 
+    /**
+     * Context key used to pass stateful decoder state across retry requests.
+     */
+    public static final String STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY
+        = "azure-storage-structured-message-decoder-state";
+
     private Constants() {
     }
 
diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java
@@ -14,16 +14,25 @@
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.storage.common.DownloadContentValidationOptions;
 import com.azure.storage.common.implementation.Constants;
-import com.azure.storage.common.implementation.structuredmessage.StructuredMessageDecodingStream;
+import com.azure.storage.common.implementation.structuredmessage.StructuredMessageDecoder;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * This is a decoding policy in an {@link com.azure.core.http.HttpPipeline} to decode structured messages in
  * storage download requests. The policy checks for a context value to determine when to apply structured message decoding.
+ *
+ * <p>The policy supports smart retries by maintaining decoder state across network interruptions, ensuring:
+ * <ul>
+ *   <li>All received segment checksums are validated before retry</li>
+ *   <li>Exact encoded and decoded byte positions are tracked</li>
+ *   <li>Decoder state is preserved across retry requests</li>
+ *   <li>Retries continue from the correct offset after network faults</li>
+ * </ul>
  */
 public class StorageContentValidationDecoderPolicy implements HttpPipelinePolicy {
     private static final ClientLogger LOGGER = new ClientLogger(StorageContentValidationDecoderPolicy.class);
@@ -51,15 +60,75 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
             Long contentLength = getContentLength(httpResponse.getHeaders());
 
             if (contentLength != null && contentLength > 0 && validationOptions != null) {
-                Flux<ByteBuffer> decodedStream = StructuredMessageDecodingStream
-                    .wrapStreamIfNeeded(httpResponse.getBody(), contentLength, validationOptions);
-                return new DecodedResponse(httpResponse, decodedStream);
+                // Get or create decoder with state tracking
+                DecoderState decoderState = getOrCreateDecoderState(context, contentLength);
+
+                // Decode using the stateful decoder
+                Flux<ByteBuffer> decodedStream = decodeStream(httpResponse.getBody(), decoderState);
+
+                // Update context with decoder state for potential retries
+                context.setData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
+
+                return new DecodedResponse(httpResponse, decodedStream, decoderState);
             }
 
             return httpResponse;
         });
     }
 
+    /**
+     * Decodes a stream of byte buffers using the decoder state.
+     *
+     * @param encodedFlux The flux of encoded byte buffers.
+     * @param state The decoder state.
+     * @return A flux of decoded byte buffers.
+     */
+    private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState state) {
+        return encodedFlux.concatMap(encodedBuffer -> {
+            try {
+                // Combine with pending data if any
+                ByteBuffer dataToProcess = state.combineWithPending(encodedBuffer);
+
+                // Track encoded bytes
+                int encodedBytesInBuffer = encodedBuffer.remaining();
+                state.totalEncodedBytesProcessed.addAndGet(encodedBytesInBuffer);
+
+                // Try to decode what we have - decoder handles partial data
+                int availableSize = dataToProcess.remaining();
+                ByteBuffer decodedData = state.decoder.decode(dataToProcess.duplicate(), availableSize);
+
+                // Track decoded bytes
+                int decodedBytes = decodedData.remaining();
+                state.totalBytesDecoded.addAndGet(decodedBytes);
+
+                // Store any remaining unprocessed data for next iteration
+                if (dataToProcess.hasRemaining()) {
+                    state.updatePendingBuffer(dataToProcess);
+                } else {
+                    state.pendingBuffer = null;
+                }
+
+                // Return decoded data if any
+                if (decodedBytes > 0) {
+                    return Flux.just(decodedData);
+                } else {
+                    return Flux.empty();
+                }
+            } catch (Exception e) {
+                LOGGER.error("Failed to decode structured message chunk: " + e.getMessage(), e);
+                return Flux.error(e);
+            }
+        }).doOnComplete(() -> {
+            // Finalize when stream completes
+            try {
+                state.decoder.finalizeDecoding();
+            } catch (IllegalArgumentException e) {
+                // Expected if we haven't received all data yet (e.g., interrupted download)
+                LOGGER.verbose("Decoding not finalized - may resume on retry: " + e.getMessage());
+            }
+        });
+    }
+
     /**
      * Checks if structured message decoding should be applied based on context.
      *
@@ -104,26 +173,131 @@ private Long getContentLength(HttpHeaders headers) {
     }
 
     /**
-     * Checks if the response is a download response (GET request with body).
+     * 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.
      *
      * @param httpResponse The HTTP response.
      * @return true if it's a download response, false otherwise.
      */
     private boolean isDownloadResponse(HttpResponse httpResponse) {
-        return httpResponse.getRequest().getHttpMethod() == HttpMethod.GET && httpResponse.getBody() != null;
+        HttpMethod method = httpResponse.getRequest().getHttpMethod();
+        return method == HttpMethod.GET && httpResponse.getStatusCode() / 100 == 2;
     }
 
     /**
-     * HTTP response wrapper that provides a decoded response body.
+     * State holder for the structured message decoder that tracks decoding progress
+     * across network interruptions.
      */
-    static class DecodedResponse extends HttpResponse {
-        private final Flux<ByteBuffer> decodedBody;
+    public static class DecoderState {
+        private final StructuredMessageDecoder decoder;
+        private final long expectedContentLength;
+        private final AtomicLong totalBytesDecoded;
+        private final AtomicLong totalEncodedBytesProcessed;
+        private ByteBuffer pendingBuffer;
+
+        /**
+         * Creates a new decoder state.
+         *
+         * @param expectedContentLength The expected length of the encoded content.
+         */
+        public DecoderState(long expectedContentLength) {
+            this.expectedContentLength = expectedContentLength;
+            this.decoder = new StructuredMessageDecoder(expectedContentLength);
+            this.totalBytesDecoded = new AtomicLong(0);
+            this.totalEncodedBytesProcessed = new AtomicLong(0);
+            this.pendingBuffer = null;
+        }
+
+        /**
+         * Combines pending buffer with new data.
+         *
+         * @param newBuffer The new buffer to combine.
+         * @return Combined buffer.
+         */
+        private ByteBuffer combineWithPending(ByteBuffer newBuffer) {
+            if (pendingBuffer == null || !pendingBuffer.hasRemaining()) {
+                return newBuffer.duplicate();
+            }
+
+            ByteBuffer combined = ByteBuffer.allocate(pendingBuffer.remaining() + newBuffer.remaining());
+            combined.put(pendingBuffer.duplicate());
+            combined.put(newBuffer.duplicate());
+            combined.flip();
+            return combined;
+        }
+
+        /**
+         * Updates the pending buffer with remaining data.
+         *
+         * @param dataToProcess The buffer with remaining data.
+         */
+        private void updatePendingBuffer(ByteBuffer dataToProcess) {
+            pendingBuffer = ByteBuffer.allocate(dataToProcess.remaining());
+            pendingBuffer.put(dataToProcess);
+            pendingBuffer.flip();
+        }
+
+        /**
+         * Gets the total number of decoded bytes processed so far.
+         *
+         * @return The total decoded bytes.
+         */
+        public long getTotalBytesDecoded() {
+            return totalBytesDecoded.get();
+        }
+
+        /**
+         * Gets the total number of encoded bytes processed so far.
+         *
+         * @return The total encoded bytes processed.
+         */
+        public long getTotalEncodedBytesProcessed() {
+            return totalEncodedBytesProcessed.get();
+        }
+
+        /**
+         * Checks if the decoder has finalized.
+         *
+         * @return true if finalized, false otherwise.
+         */
+        public boolean isFinalized() {
+            return totalEncodedBytesProcessed.get() >= expectedContentLength;
+        }
+    }
+
+    /**
+     * Decoded HTTP response that wraps the original response with a decoded stream.
+     */
+    private static class DecodedResponse extends HttpResponse {
         private final HttpResponse originalResponse;
+        private final Flux<ByteBuffer> decodedBody;
+        private final DecoderState decoderState;
 
-        DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody) {
-            super(httpResponse.getRequest());
-            this.originalResponse = httpResponse;
+        /**
+         * Creates a new decoded response.
+         *
+         * @param originalResponse The original HTTP response.
+         * @param decodedBody The decoded body stream.
+         * @param decoderState The decoder state.
+         */
+        DecodedResponse(HttpResponse originalResponse, Flux<ByteBuffer> decodedBody, DecoderState decoderState) {
+            super(originalResponse.getRequest());
+            this.originalResponse = originalResponse;
             this.decodedBody = decodedBody;
+            this.decoderState = decoderState;
         }
 
         @Override
@@ -153,12 +327,21 @@ public Mono<byte[]> getBodyAsByteArray() {
 
         @Override
         public Mono<String> getBodyAsString() {
-            return getBodyAsByteArray().map(String::new);
+            return getBodyAsByteArray().map(bytes -> new String(bytes, Charset.defaultCharset()));
         }
 
         @Override
         public Mono<String> getBodyAsString(Charset charset) {
             return getBodyAsByteArray().map(bytes -> new String(bytes, charset));
         }
+
+        /**
+         * Gets the decoder state.
+         *
+         * @return The decoder state.
+         */
+        public DecoderState getDecoderState() {
+            return decoderState;
+        }
     }
 }
