smart retry changes

Author: gunjansingh-msft

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

@@ -114,6 +114,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
@@ -1318,6 +1319,18 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                     finalCount = finalRange.getCount();
                 }
 
+                AtomicReference<StorageContentValidationDecoderPolicy.DecoderState> decoderStateRef
+                    = new AtomicReference<>();
+                if (contentValidationOptions != null
+                    && contentValidationOptions.isStructuredMessageValidationEnabled()) {
+                    Object decoderStateObj
+                        = firstRangeContext.getData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY)
+                            .orElse(null);
+                    if (decoderStateObj instanceof StorageContentValidationDecoderPolicy.DecoderState) {
+                        decoderStateRef.set((StorageContentValidationDecoderPolicy.DecoderState) decoderStateObj);
+                    }
+                }
+
                 // The resume function takes throwable and offset at the destination.
                 // I.e. offset is relative to the starting point.
                 BiFunction<Throwable, Long, Mono<StreamResponse>> onDownloadErrorResume = (throwable, offset) -> {
@@ -1326,18 +1339,28 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                     }
 
                     long newCount = finalCount - offset;
+                    StorageContentValidationDecoderPolicy.DecoderState decoderState = null;
+                    long expectedEncodedLength = finalCount;
+                    long encodedProgress = offset;
+
+                    if (contentValidationOptions != null
+                        && contentValidationOptions.isStructuredMessageValidationEnabled()) {
+                        decoderState = decoderStateRef.get();
+
+                        if (decoderState == null) {
+                            Object decoderStateObj
+                                = firstRangeContext.getData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY)
+                                    .orElse(null);
 
-                    /*
-                     * It's possible that the network stream will throw an error after emitting all data but before
-                     * completing. Issuing a retry at this stage would leave the download in a bad state with
-                     * incorrect count and offset values. Because we have read the intended amount of data, we can
-                     * ignore the error at the end of the stream.
-                     */
-                    if (newCount == 0) {
-                        LOGGER.warning("Exception encountered in ReliableDownload after all data read from the network "
-                            + "but before stream signaled completion. Returning success as all data was downloaded. "
-                            + "Exception message: " + throwable.getMessage());
-                        return Mono.empty();
+                            if (decoderStateObj instanceof StorageContentValidationDecoderPolicy.DecoderState) {
+                                decoderState = (StorageContentValidationDecoderPolicy.DecoderState) decoderStateObj;
+                            }
+                        }
+
+                        if (decoderState != null) {
+                            expectedEncodedLength = decoderState.getExpectedContentLength();
+                            encodedProgress = decoderState.getTotalEncodedBytesProcessed();
+                        }
                     }
 
                     try {
@@ -1349,53 +1372,38 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                         // based on the encoded bytes processed, not the decoded bytes
                         if (contentValidationOptions != null
                             && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                            // 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);
-
-                            if (decoderStateObj instanceof StorageContentValidationDecoderPolicy.DecoderState) {
-                                StorageContentValidationDecoderPolicy.DecoderState decoderState
-                                    = (StorageContentValidationDecoderPolicy.DecoderState) decoderStateObj;
+                            long retryStartOffset = -1;
 
-                                // Use getRetryOffset() to get the correct offset for retry
-                                // This accounts for pending bytes that have been received but not yet consumed
-                                long encodedOffset = decoderState.getRetryOffset();
-                                long remainingCount = finalCount - encodedOffset;
-                                retryRange = new BlobRange(initialOffset + encodedOffset, remainingCount);
+                            // First try to use decoder state (authoritative)
+                            if (decoderState != null) {
+                                // Always rewind decoder to last validated boundary before retrying.
+                                retryStartOffset = decoderState.resetForRetry();
 
-                                LOGGER.info(
-                                    "Structured message smart retry: resuming from offset {} (initial={}, encoded={})",
-                                    initialOffset + encodedOffset, initialOffset, encodedOffset);
-
-                                // Preserve the decoder state for the retry
                                 retryContext = retryContext
                                     .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
-                            } else {
-                                // No decoder state available, try to parse retry offset from exception message
-                                // The exception message contains RETRY-START-OFFSET=<number> token
-                                long retryStartOffset = StorageContentValidationDecoderPolicy
+                                decoderStateRef.set(decoderState);
+                            }
+
+                            // If no decoder state or no retry offset from state, fall back to parsed token or offset.
+                            if (retryStartOffset < 0) {
+                                retryStartOffset = StorageContentValidationDecoderPolicy
                                     .parseRetryStartOffset(throwable.getMessage());
-                                if (retryStartOffset >= 0) {
-                                    long remainingCount = finalCount - retryStartOffset;
-                                    // Validate remainingCount to avoid negative values
-                                    if (remainingCount <= 0) {
-                                        LOGGER.warning("Retry offset {} exceeds finalCount {}, using fallback",
-                                            retryStartOffset, finalCount);
-                                        retryRange = new BlobRange(initialOffset + offset, newCount);
-                                    } else {
-                                        retryRange = new BlobRange(initialOffset + retryStartOffset, remainingCount);
-
-                                        LOGGER.info(
-                                            "Structured message smart retry from exception: resuming from offset {} "
-                                                + "(initial={}, parsed={})",
-                                            initialOffset + retryStartOffset, initialOffset, retryStartOffset);
-                                    }
-                                } else {
-                                    // Fallback to normal retry logic if no offset found
-                                    retryRange = new BlobRange(initialOffset + offset, newCount);
-                                }
                             }
+                            if (retryStartOffset < 0) {
+                                retryStartOffset = offset;
+                            }
+
+                            long remainingCount = expectedEncodedLength - retryStartOffset;
+                            if (remainingCount < 0) {
+                                remainingCount = expectedEncodedLength - offset;
+                                retryStartOffset = offset;
+                            }
+
+                            retryRange = new BlobRange(initialOffset + retryStartOffset, remainingCount);
+
+                            LOGGER.info(
+                                "Structured message smart retry: resuming from offset {} (initial={}, encoded={}, remaining={})",
+                                initialOffset + retryStartOffset, initialOffset, retryStartOffset, remainingCount);
                         } else {
                             // For non-structured downloads, use smart retry from the interrupted offset
                             retryRange = new BlobRange(initialOffset + offset, newCount);
@@ -1410,6 +1418,7 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                 // Structured message decoding is now handled by StructuredMessageDecoderPolicy
                 return BlobDownloadAsyncResponseConstructorProxy.create(response, onDownloadErrorResume, finalOptions);
             });
+
     }
 
     Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, DownloadRetryOptions options,

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

@@ -5,6 +5,7 @@
 
 import com.azure.core.test.utils.TestUtils;
 import com.azure.core.util.FluxUtil;
+import com.azure.core.util.Context;
 import com.azure.storage.blob.models.BlobRange;
 import com.azure.storage.blob.models.BlobRequestConditions;
 import com.azure.storage.blob.models.DownloadRetryOptions;
@@ -20,6 +21,7 @@
 import reactor.test.StepVerifier;
 
 import java.io.IOException;
+import java.io.ByteArrayOutputStream;
 import java.nio.ByteBuffer;
 import java.util.List;
 
@@ -285,6 +287,83 @@ public void downloadStreamWithResponseContentValidationSmartRetry() throws IOExc
         }
     }
 
+    @Test
+    public void downloadStreamWithResponseContentValidationSmartRetryVariousSizes() throws IOException {
+        int[] dataSizes = new int[] { Constants.KB, 1500, 3 * Constants.KB + 128 };
+        int[] segmentSizes = new int[] { 512, 700, 1024 };
+
+        for (int i = 0; i < dataSizes.length; i++) {
+            byte[] randomData = getRandomByteArray(dataSizes[i]);
+            int segmentSize = segmentSizes[i % segmentSizes.length];
+
+            StructuredMessageEncoder encoder
+                = new StructuredMessageEncoder(randomData.length, segmentSize, StructuredMessageFlags.STORAGE_CRC64);
+            ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+            Flux<ByteBuffer> input = Flux.just(encodedData);
+
+            // Create a policy that will simulate 2 network interruptions for each run
+            MockPartialResponsePolicy mockPolicy = new MockPartialResponsePolicy(2);
+
+            // Upload the encoded data
+            bc.upload(input, null, true).block();
+
+            // Create a download client with both the mock policy AND the decoder policy
+            StorageContentValidationDecoderPolicy decoderPolicy = new StorageContentValidationDecoderPolicy();
+            BlobAsyncClient downloadClient = getBlobAsyncClient(ENVIRONMENT.getPrimaryAccount().getCredential(),
+                bc.getBlobUrl(), mockPolicy, decoderPolicy);
+
+            DownloadContentValidationOptions validationOptions
+                = new DownloadContentValidationOptions().setStructuredMessageValidationEnabled(true);
+            DownloadRetryOptions retryOptions = new DownloadRetryOptions().setMaxRetryRequests(5);
+
+            StepVerifier.create(downloadClient
+                .downloadStreamWithResponse((BlobRange) null, retryOptions, (BlobRequestConditions) null, false,
+                    validationOptions)
+                .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue()))).assertNext(r -> {
+                    TestUtils.assertArraysEqual(r, randomData);
+                }).verifyComplete();
+
+            assertEquals(0, mockPolicy.getTriesRemaining());
+            List<String> rangeHeaders = mockPolicy.getRangeHeaders();
+            assertTrue(rangeHeaders.size() > 0, "Expected range headers for retries");
+            assertTrue(rangeHeaders.get(0).startsWith("bytes=0-"), "First request should start from offset 0");
+        }
+    }
+
+    @Test
+    public void downloadStreamWithResponseContentValidationSmartRetrySync() throws IOException {
+        byte[] randomData = getRandomByteArray(Constants.KB);
+        StructuredMessageEncoder encoder
+            = new StructuredMessageEncoder(randomData.length, 512, StructuredMessageFlags.STORAGE_CRC64);
+        ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+        Flux<ByteBuffer> input = Flux.just(encodedData);
+        bc.upload(input, null, true).block();
+
+        MockPartialResponsePolicy mockPolicy = new MockPartialResponsePolicy(3);
+        StorageContentValidationDecoderPolicy decoderPolicy = new StorageContentValidationDecoderPolicy();
+        BlobClient downloadClient = getBlobClient(ENVIRONMENT.getPrimaryAccount().getCredential(), bc.getBlobUrl(),
+            mockPolicy, decoderPolicy);
+
+        DownloadContentValidationOptions validationOptions
+            = new DownloadContentValidationOptions().setStructuredMessageValidationEnabled(true);
+        DownloadRetryOptions retryOptions = new DownloadRetryOptions().setMaxRetryRequests(5);
+
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        Context ctx = new Context(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true)
+            .addData(Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY, validationOptions);
+
+        downloadClient.downloadStreamWithResponse(out, null, retryOptions, null, false, null, ctx);
+
+        TestUtils.assertArraysEqual(out.toByteArray(), randomData);
+
+        assertEquals(0, mockPolicy.getTriesRemaining());
+        List<String> rangeHeaders = mockPolicy.getRangeHeaders();
+        assertTrue(rangeHeaders.size() > 0, "Expected range headers for retries");
+        assertTrue(rangeHeaders.get(0).startsWith("bytes=0-"), "First request should start from offset 0");
+    }
+
     @Test
     public void downloadStreamWithResponseContentValidationSmartRetryMultipleSegments() throws IOException {
         // Test smart retry with multiple segments to ensure checksum validation

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

@@ -332,6 +332,20 @@ public boolean isHeaderRead() {
         return messageLength != -1;
     }
 
+    /**
+     * Gets the number of encoded bytes that have been seen but not yet fully
+     * processed by the decoder (pending bytes).
+     *
+     * <p>This is used by smart-retry logic to determine the absolute encoded
+     * offset that a retry request should start from while still preserving
+     * the decoder's buffered state.</p>
+     *
+     * @return The number of pending encoded bytes buffered by the decoder.
+     */
+    public int getPendingEncodedByteCount() {
+        return pendingBytes.size();
+    }
+
     /**
      * Reads the message header if we have enough bytes.
      *
@@ -611,6 +625,13 @@ public DecodeResult decodeChunk(ByteBuffer buffer) {
 
             // Step 2: Process segments
             while (messageOffset < messageLength) {
+                // If all segments are done, proceed to message footer before attempting any new segment header.
+                if (currentSegmentNumber == numSegments && currentSegmentContentOffset == currentSegmentContentLength) {
+                    if (!tryReadMessageFooter(buffer)) {
+                        break;
+                    }
+                }
+
                 // Read segment header if needed
                 if (currentSegmentContentOffset == currentSegmentContentLength) {
                     if (!tryReadSegmentHeader(buffer)) {