smart retry changes

Author: gunjansingh-msft

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/implementation/util/BuilderHelper.java

@@ -141,7 +141,12 @@ public static HttpPipeline buildPipeline(StorageSharedKeyCredential storageShare
 
         HttpPolicyProviders.addAfterRetryPolicies(policies);
 
-        policies.add(new StorageContentValidationDecoderPolicy());
+        // Only add the structured message decoder once; allow callers to inject their own for ordering with test
+        // policies (e.g., fault injection before decoding).
+        boolean hasDecoder = policies.stream().anyMatch(p -> p instanceof StorageContentValidationDecoderPolicy);
+        if (!hasDecoder) {
+            policies.add(new StorageContentValidationDecoderPolicy());
+        }
 
         policies.add(getResponseValidationPolicy());
 

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

@@ -1290,11 +1290,14 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
             ? new Context("azure-eagerly-convert-headers", true)
             : context.addData("azure-eagerly-convert-headers", true);
 
+        AtomicReference<StorageContentValidationDecoderPolicy.DecoderState> decoderStateRef = new AtomicReference<>();
+
         // Add structured message decoding context if enabled
         final Context firstRangeContext;
         if (contentValidationOptions != null && contentValidationOptions.isStructuredMessageValidationEnabled()) {
             firstRangeContext = initialContext.addData(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true)
-                .addData(Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY, contentValidationOptions);
+                .addData(Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY, contentValidationOptions)
+                .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_REF_CONTEXT_KEY, decoderStateRef);
         } else {
             firstRangeContext = initialContext;
         }
@@ -1319,18 +1322,6 @@ 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) -> {
@@ -1342,6 +1333,8 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                     StorageContentValidationDecoderPolicy.DecoderState decoderState = null;
                     long expectedEncodedLength = finalCount;
                     long encodedProgress = offset;
+                    long retryStartOffset = -1;
+                    boolean noBytesEmitted = offset == 0;
 
                     if (contentValidationOptions != null
                         && contentValidationOptions.isStructuredMessageValidationEnabled()) {
@@ -1357,12 +1350,37 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                             }
                         }
 
+                        // If the decoder has already finalized, discard it and restart from the beginning.
+                        if (decoderState != null && decoderState.isFinalized()) {
+                            decoderState = null;
+                        }
+
                         if (decoderState != null) {
                             expectedEncodedLength = decoderState.getExpectedContentLength();
                             encodedProgress = decoderState.getTotalEncodedBytesProcessed();
+
+                            long boundaryDecoded = decoderState.getDecodedBytesAtLastCompleteSegment();
+                            if (offset < boundaryDecoded || noBytesEmitted) {
+                                // We haven't emitted the bytes represented by this decoder boundary; restart clean.
+                                decoderState = null;
+                            } else {
+                                long bytesToSkip = Math.max(0, offset - boundaryDecoded);
+                                decoderState.setDecodedBytesToSkip(bytesToSkip);
+
+                                // Always rewind decoder to last validated boundary before retrying.
+                                retryStartOffset = decoderState.resetForRetry();
+                            }
                         }
                     }
 
+                    if (decoderState == null) {
+                        // No decoder state available (likely failed before policy captured it) or no bytes emitted;
+                        // restart from beginning to avoid skipping data.
+                        retryStartOffset = noBytesEmitted ? 0 : -1;
+                        encodedProgress = noBytesEmitted ? 0 : encodedProgress;
+                        decoderStateRef.set(null);
+                    }
+
                     try {
                         // For retry context, preserve decoder state if structured message validation is enabled
                         Context retryContext = firstRangeContext;
@@ -1372,25 +1390,23 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                         // based on the encoded bytes processed, not the decoded bytes
                         if (contentValidationOptions != null
                             && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                            long retryStartOffset = -1;
-
-                            // First try to use decoder state (authoritative)
-                            if (decoderState != null) {
-                                // Always rewind decoder to last validated boundary before retrying.
-                                retryStartOffset = decoderState.resetForRetry();
-
-                                retryContext = retryContext
-                                    .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
-                                decoderStateRef.set(decoderState);
-                            }
-
                             // If no decoder state or no retry offset from state, fall back to parsed token or offset.
-                            if (retryStartOffset < 0) {
+                            if (retryStartOffset < 0 && !noBytesEmitted) {
                                 retryStartOffset = StorageContentValidationDecoderPolicy
                                     .parseRetryStartOffset(throwable.getMessage());
                             }
                             if (retryStartOffset < 0) {
-                                retryStartOffset = offset;
+                                retryStartOffset = noBytesEmitted ? 0 : offset;
+                            }
+
+                            if (decoderState != null) {
+                                retryContext = retryContext
+                                    .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, decoderState);
+                                decoderStateRef.set(decoderState);
+                            } else {
+                                // Ensure we don't carry a stale decoder state into a full restart.
+                                retryContext = retryContext
+                                    .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_CONTEXT_KEY, null);
                             }
 
                             long remainingCount = expectedEncodedLength - retryStartOffset;