simplifying retry mechanism

Author: gunjansingh-msft

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/implementation/accesshelpers/BlobDownloadAsyncResponseConstructorProxy.java

@@ -9,10 +9,8 @@
 import com.azure.core.http.rest.StreamResponse;
 import com.azure.storage.blob.models.BlobDownloadAsyncResponse;
 import com.azure.storage.blob.models.DownloadRetryOptions;
-import com.azure.storage.common.policy.DecoderState;
 import reactor.core.publisher.Mono;
 
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 
 /**
@@ -37,16 +35,7 @@ public interface BlobDownloadAsyncResponseConstructorAccessor {
          * @param retryOptions Retry options.
          */
         BlobDownloadAsyncResponse create(StreamResponse sourceResponse,
-            BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-            AtomicReference<DecoderState> decoderStateRef);
-
-        /**
-         * Gets the current decoder state from a {@link BlobDownloadAsyncResponse}.
-         *
-         * @param response The {@link BlobDownloadAsyncResponse}.
-         * @return The current decoder state, or null if not available.
-         */
-        DecoderState getDecoderState(BlobDownloadAsyncResponse response);
+            BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions);
     }
 
     /**
@@ -67,8 +56,7 @@ public static void setAccessor(
      * @param retryOptions Retry options.
      */
     public static BlobDownloadAsyncResponse create(StreamResponse sourceResponse,
-        BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-        AtomicReference<DecoderState> decoderStateRef) {
+        BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions) {
         // This looks odd but is necessary, it is possible to engage the access helper before anywhere else in the
         // application accesses BlobDownloadAsyncResponse which triggers the accessor to be configured. So, if the accessor
         // is null this effectively pokes the class to set up the accessor.
@@ -78,16 +66,6 @@ public static BlobDownloadAsyncResponse create(StreamResponse sourceResponse,
         }
 
         assert accessor != null;
-        return accessor.create(sourceResponse, onErrorResume, retryOptions, decoderStateRef);
-    }
-
-    /**
-     * Gets the current decoder state from a {@link BlobDownloadAsyncResponse}.
-     *
-     * @param response The {@link BlobDownloadAsyncResponse}.
-     * @return The decoder state, or null if not available.
-     */
-    public static DecoderState getDecoderState(BlobDownloadAsyncResponse response) {
-        return accessor.getDecoderState(response);
+        return accessor.create(sourceResponse, onErrorResume, retryOptions);
     }
 }

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

@@ -6,11 +6,6 @@
 import com.azure.core.util.Context;
 import com.azure.storage.common.StorageChecksumAlgorithm;
 import com.azure.storage.common.implementation.Constants;
-import com.azure.storage.common.policy.AggregateCrcState;
-
-import java.util.concurrent.atomic.AtomicReference;
-
-import com.azure.storage.common.policy.DecoderState;
 
 /**
  * Centralizes download content validation decisions based on {@link StorageChecksumAlgorithm}.
@@ -39,23 +34,17 @@ public static StorageChecksumAlgorithm resolveAlgorithm(StorageChecksumAlgorithm
     }
 
     /**
-     * Adds structured message decoding context keys when CRC64/AUTO validation is active.
+     * Adds structured message decoding context key when CRC64/AUTO validation is active.
      *
      * @param context The base context to augment. Null is treated as {@link Context#NONE}.
      * @param algorithm The resolved checksum algorithm.
-     * @param decoderStateRef Holder for decoder state, populated by the policy.
-     * @param aggregateCrcState CRC aggregation state across retries.
      * @return The augmented context.
      */
-    public static Context applyStructuredMessageContext(Context context, StorageChecksumAlgorithm algorithm,
-        AtomicReference<DecoderState> decoderStateRef, AggregateCrcState aggregateCrcState) {
+    public static Context applyStructuredMessageContext(Context context, StorageChecksumAlgorithm algorithm) {
         Context base = context == null ? Context.NONE : context;
         if (!isStructuredMessageAlgorithm(algorithm)) {
             return base;
         }
-        return base.addData(Constants.STRUCTURED_MESSAGE_RESPONSE_SCOPED_CONTEXT_KEY, true)
-            .addData(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true)
-            .addData(Constants.STRUCTURED_MESSAGE_DECODER_STATE_REF_CONTEXT_KEY, decoderStateRef)
-            .addData(Constants.STRUCTURED_MESSAGE_AGGREGATE_CRC_CONTEXT_KEY, aggregateCrcState);
+        return base.addData(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true);
     }
 }

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/models/BlobDownloadAsyncResponse.java

@@ -14,17 +14,14 @@
 import com.azure.storage.blob.implementation.models.BlobsDownloadHeaders;
 import com.azure.storage.blob.implementation.util.ModelHelper;
 import com.azure.core.util.logging.ClientLogger;
-import com.azure.storage.common.policy.DecoderState;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
 import java.nio.channels.AsynchronousByteChannel;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 
 /**
@@ -40,15 +37,9 @@ public final class BlobDownloadAsyncResponse extends ResponseBase<BlobDownloadHe
             .setAccessor(new BlobDownloadAsyncResponseConstructorProxy.BlobDownloadAsyncResponseConstructorAccessor() {
                 @Override
                 public BlobDownloadAsyncResponse create(StreamResponse sourceResponse,
-                    BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-                    AtomicReference<DecoderState> decoderStateRef) {
-                    return new BlobDownloadAsyncResponse(sourceResponse, onErrorResume, retryOptions, decoderStateRef);
-                }
-
-                @Override
-                public DecoderState getDecoderState(BlobDownloadAsyncResponse response) {
-                    AtomicReference<DecoderState> ref = response.decoderStateRef;
-                    return ref == null ? null : ref.get();
+                    BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume,
+                    DownloadRetryOptions retryOptions) {
+                    return new BlobDownloadAsyncResponse(sourceResponse, onErrorResume, retryOptions);
                 }
             });
     }
@@ -58,7 +49,6 @@ public DecoderState getDecoderState(BlobDownloadAsyncResponse response) {
     private final StreamResponse sourceResponse;
     private final BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume;
     private final DownloadRetryOptions retryOptions;
-    private final AtomicReference<DecoderState> decoderStateRef;
 
     /**
      * Constructs a {@link BlobDownloadAsyncResponse}.
@@ -75,38 +65,15 @@ public BlobDownloadAsyncResponse(HttpRequest request, int statusCode, HttpHeader
         this.sourceResponse = null;
         this.onErrorResume = null;
         this.retryOptions = null;
-        this.decoderStateRef = null;
     }
 
-    /**
-     * Constructs a {@link BlobDownloadAsyncResponse}.
-     *
-     * @param sourceResponse The initial Stream Response
-     * @param onErrorResume Function used to resume.
-     * @param retryOptions Retry options.
-     */
     BlobDownloadAsyncResponse(StreamResponse sourceResponse,
         BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions) {
-        this(sourceResponse, onErrorResume, retryOptions, null);
-    }
-
-    BlobDownloadAsyncResponse(StreamResponse sourceResponse,
-        BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-        AtomicReference<DecoderState> decoderStateRef) {
-        this(sourceResponse, onErrorResume, retryOptions, decoderStateRef, extractHeaders(sourceResponse));
-    }
-
-    private BlobDownloadAsyncResponse(StreamResponse sourceResponse,
-        BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-        AtomicReference<DecoderState> decoderStateRef, BlobDownloadHeaders deserializedHeaders) {
         super(sourceResponse.getRequest(), sourceResponse.getStatusCode(), sourceResponse.getHeaders(),
-            createResponseFluxWithContentCrc(sourceResponse, onErrorResume, retryOptions, decoderStateRef,
-                deserializedHeaders),
-            deserializedHeaders);
+            createResponseFlux(sourceResponse, onErrorResume, retryOptions), extractHeaders(sourceResponse));
         this.sourceResponse = Objects.requireNonNull(sourceResponse, "'sourceResponse' must not be null");
         this.onErrorResume = Objects.requireNonNull(onErrorResume, "'onErrorResume' must not be null");
         this.retryOptions = Objects.requireNonNull(retryOptions, "'retryOptions' must not be null");
-        this.decoderStateRef = decoderStateRef;
     }
 
     private static BlobDownloadHeaders extractHeaders(StreamResponse response) {
@@ -124,28 +91,6 @@ private static Flux<ByteBuffer> createResponseFlux(StreamResponse sourceResponse
             .defaultIfEmpty(EMPTY_BUFFER);
     }
 
-    /**
-     * Builds the response flux and populates ContentCrc64 on the deserialized headers when structured message
-     * decoding completes successfully.
-     */
-    private static Flux<ByteBuffer> createResponseFluxWithContentCrc(StreamResponse sourceResponse,
-        BiFunction<Throwable, Long, Mono<StreamResponse>> onErrorResume, DownloadRetryOptions retryOptions,
-        AtomicReference<DecoderState> decoderStateRef, BlobDownloadHeaders deserializedHeaders) {
-        Flux<ByteBuffer> flux = createResponseFlux(sourceResponse, onErrorResume, retryOptions);
-        if (decoderStateRef != null && deserializedHeaders != null) {
-            flux = flux.doOnComplete(() -> {
-                DecoderState state = decoderStateRef.get();
-                if (state != null && state.isFinalized()) {
-                    long crc = state.getComposedCrc64();
-                    byte[] crcBytes = new byte[8];
-                    ByteBuffer.wrap(crcBytes).order(ByteOrder.LITTLE_ENDIAN).putLong(crc);
-                    deserializedHeaders.setContentCrc64(crcBytes);
-                }
-            });
-        }
-        return flux;
-    }
-
     /**
      * Transfers content bytes to the {@link AsynchronousByteChannel}.
      * @param channel The destination {@link AsynchronousByteChannel}.