5945 retrypolicy response headers

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

@@ -15,6 +15,17 @@
  * A collection of headers on an HTTP request or response.
  */
 public class HttpHeaders implements Iterable<HttpHeader> {
+
+    /**
+     * The header which will have retry after value in milliseconds.
+     */
+    public static final String AZURE_RETRY_AFTER_MS_HEADER = "retry-after-ms";
+
+    /**
+     * The header which will have retry after value in milliseconds.
+     */
+    public static final String AZURE_X_MS_RETRY_AFTER_MS_HEADER = "x-ms-retry-after-ms";
+
     private final Map<String, HttpHeader> headers = new ConcurrentHashMap<>();
 
     /**

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

@@ -5,12 +5,17 @@
 
 import static com.azure.core.util.CoreUtils.isNullOrEmpty;
 
+import java.time.temporal.ChronoUnit;
+
+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.util.logging.ClientLogger;
+
 import reactor.core.publisher.Mono;
 
 import java.time.Duration;
@@ -20,25 +25,58 @@
  */
 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 String retryAfterHeader;
+    private final ChronoUnit retryAfterTimeUnit;
+
     /**
      * Creates a default {@link ExponentialBackoff} retry policy.
      */
     public RetryPolicy() {
-        this(new ExponentialBackoff());
+        this(new ExponentialBackoff(), null, null);
+    }
+
+    /**
+     * Creates a default {@link ExponentialBackoff} retry policy.
+     * @param retryAfterHeader The retry after http header name to be used get retry after value from
+     * @param retryAfterTimeUnit The time unit to use while applying retry based on value specified in
+     *      * {@code retryAfterHeader} in {@link HttpResponse}.
+     * {@link HttpResponse}.
+     */
+    public RetryPolicy(String retryAfterHeader, ChronoUnit retryAfterTimeUnit) {
+        this(new ExponentialBackoff(), retryAfterHeader, retryAfterTimeUnit);
     }
 
     /**
      * 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 = Objects.requireNonNull(retryStrategy, "'retryStrategy' cannot be null");
+        this(retryStrategy, null, null);
+    }
+
+    /**
+     * Creates a {@link RetryPolicy} with the provided {@link RetryStrategy} and {@code retryAfterHeader}.
+     *
+     * @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}. The value {@code null} is valid.
+     * @param retryAfterTimeUnit The time unit to use while applying retry based on value specified in
+     * {@code retryAfterHeader} in {@link HttpResponse}.The value {@code null} is valid only in case when
+     * {@code retryAfterHeader} is empty or {@code null}.
+     * @throws NullPointerException if {@code retryStrategy} is {@code null}.
+     */
+    public RetryPolicy(RetryStrategy retryStrategy, String retryAfterHeader, ChronoUnit retryAfterTimeUnit) {
+        this.retryStrategy = Objects.requireNonNull(retryStrategy, "'retryStrategy' cannot be null.");
+        this.retryAfterHeader = retryAfterHeader;
+        this.retryAfterTimeUnit = retryAfterTimeUnit;
+        if (!isNullOrEmpty(retryAfterHeader)) {
+            Objects.requireNonNull(retryAfterTimeUnit, "'retryAfterTimeUnit' cannot be null.");
+        }
     }
 
     @Override
@@ -94,14 +132,18 @@ private Duration determineDelayDuration(HttpResponse response, int tryCount) {
             return retryStrategy.calculateRetryDelay(tryCount);
         }
 
-        String retryHeader = response.getHeaderValue(RETRY_AFTER_MS_HEADER);
+        String retryHeaderValue = null;
+
+        if (!isNullOrEmpty(retryAfterHeader)) {
+            retryHeaderValue = response.getHeaderValue(retryAfterHeader);
+        }
 
         // 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.of(Integer.parseInt(retryHeaderValue), retryAfterTimeUnit);
     }
 }