5945 retrypolicy response headers

sdk/core/azure-core/src/main/java/com/azure/core/http/HttpHeaderName.java

@@ -0,0 +1,54 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http;
+
+import java.util.Locale;
+
+/**
+ * Declare various headers which can be used to interact with Azure services.
+ */
+public enum HttpHeaderName {
+
+    /**
+     * The header which will have retry after value in milliseconds.
+     */
+    AZURE_RETRY_AFTER_MS_HEADER("retry-after-ms"),
+
+    /**
+     * The header which will have retry after value in milliseconds.
+     */
+     AZURE_X_MS_RETRY_AFTER_MS_HEADER("x-ms-retry-after-ms");
+
+    private final String value;
+
+    HttpHeaderName(final String value) {
+        this.value = value;
+    }
+
+    /**
+     * Returns string representation of the {@link HttpHeaderName}.
+     */
+    @Override
+    public String toString() {
+        return this.value;
+    }
+
+    /**
+     * Creates a {@link HttpHeaderName} from its string value.
+     *
+     * @param value The string value of the  {@link HttpHeaderName}.
+     * @return The  {@link HttpHeaderName} represented by the value.
+     * @throws IllegalArgumentException If a {@link HttpHeaderName} cannot be parsed from the string value.
+     */
+    public static HttpHeaderName fromString(final String value) {
+        for (HttpHeaderName httpHeaderType : values()) {
+            if (httpHeaderType.value.equalsIgnoreCase(value)) {
+                return httpHeaderType;
+            }
+        }
+
+        throw new IllegalArgumentException(String.format(Locale.US, "Could not convert %s to a HttpHeaderName",
+            value));
+    }
+}

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RetryPolicy.java

@@ -5,12 +5,16 @@
 
 import static com.azure.core.util.CoreUtils.isNullOrEmpty;
 
+import java.util.Objects;
+
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
-import java.util.Objects;
+import com.azure.core.http.HttpHeaderName;
+
 import com.azure.core.util.logging.ClientLogger;
+
 import reactor.core.publisher.Mono;
 
 import java.time.Duration;
@@ -20,25 +24,49 @@
  */
 public class RetryPolicy implements HttpPipelinePolicy {
 
-    private static final String RETRY_AFTER_MS_HEADER = "retry-after-ms";
-
     private final ClientLogger logger = new ClientLogger(RetryPolicy.class);
     private final RetryStrategy retryStrategy;
 
+    private final HttpHeaderName retryAfterHeader;
+
     /**
      * Creates a default {@link ExponentialBackoff} retry policy.
      */
     public RetryPolicy() {
-        this(new ExponentialBackoff());
+        this(new ExponentialBackoff(), HttpHeaderName.AZURE_RETRY_AFTER_MS_HEADER);
+    }
+
+    /**
+     * Creates a default {@link ExponentialBackoff} retry policy.
+     * @param retryAfterHeader The retry after http header name to be used get retry after value from
+     * {@link HttpResponse}.
+     * @throws NullPointerException if {@code retryAfterHeader} is {@code null}.
+     */
+    public RetryPolicy(HttpHeaderName retryAfterHeader) {
+        this(new ExponentialBackoff(), retryAfterHeader);
     }
 
     /**
      * Creates a RetryPolicy with the provided {@link RetryStrategy}.
      *
      * @param retryStrategy The {@link RetryStrategy} used for retries.
+     * @throws NullPointerException if {@code retryStrategy} is {@code null}.
      */
     public RetryPolicy(RetryStrategy retryStrategy) {
+        this(retryStrategy, HttpHeaderName.AZURE_RETRY_AFTER_MS_HEADER);
+    }
+
+    /**
+     * Creates a {@link RetryPolicy} with the provided {@link RetryStrategy} and {@link HttpHeaderName}.
+     *
+     * @param retryStrategy The {@link RetryStrategy} used for retries.
+     * @param retryAfterHeader The retry after http header name to be used get retry after value from
+     * {@link HttpResponse}.
+     * @throws NullPointerException if {@code retryStrategy} or {@code retryAfterHeader} is {@code null}.
+     */
+    public RetryPolicy(RetryStrategy retryStrategy, HttpHeaderName retryAfterHeader) {
         this.retryStrategy = Objects.requireNonNull(retryStrategy, "'retryStrategy' cannot be null");
+        this.retryAfterHeader = Objects.requireNonNull(retryAfterHeader, "'retryAfterHeader' cannot be null");
     }
 
     @Override
@@ -94,14 +122,14 @@ private Duration determineDelayDuration(HttpResponse response, int tryCount) {
             return retryStrategy.calculateRetryDelay(tryCount);
         }
 
-        String retryHeader = response.getHeaderValue(RETRY_AFTER_MS_HEADER);
+        String retryHeaderValue = response.getHeaderValue(retryAfterHeader.name());
 
         // Retry header is missing or empty, return the default delay duration.
-        if (isNullOrEmpty(retryHeader)) {
+        if (isNullOrEmpty(retryHeaderValue)) {
             return retryStrategy.calculateRetryDelay(tryCount);
         }
 
         // Use the response delay duration, the server returned it for a reason.
-        return Duration.ofMillis(Integer.parseInt(retryHeader));
+        return Duration.ofMillis(Integer.parseInt(retryHeaderValue));
     }
 }