5945 retrypolicy response headers

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,67 @@
  */
 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.
+     * Creates a default {@link ExponentialBackoff} retry policy. It will not use {@code retryAfterHeader}
+     * in {@link HttpResponse}.
      */
     public RetryPolicy() {
-        this(new ExponentialBackoff());
+        this(new ExponentialBackoff(), null, null);
+    }
+
+    /**
+     * Creates a default {@link ExponentialBackoff} retry policy along with provided {@code retryAfterHeader} and
+     * {@code retryAfterTimeUnit}.
+     *
+     * @param retryAfterHeader The 'retry-after' HTTP header name to lookup for the retry duration.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 Only if {@code retryAfterTimeUnit} is {@code null} and {@code retryAfterHeader}
+     * is not {@code null}.
+     */
+    public RetryPolicy(String retryAfterHeader, ChronoUnit retryAfterTimeUnit) {
+        this(new ExponentialBackoff(), retryAfterHeader, retryAfterTimeUnit);
     }
 
     /**
-     * Creates a RetryPolicy with the provided {@link RetryStrategy}.
+     * Creates a RetryPolicy with the provided {@link RetryStrategy}. It will not use {@code retryAfterHeader}
+     * in {@link HttpResponse}.
      *
      * @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}, {@code retryAfterHeader}  and
+     * {@code retryAfterTimeUnit}.
+     *
+     * @param retryStrategy The {@link RetryStrategy} used for retries.
+     * @param retryAfterHeader The 'retry-after' HTTP header name to lookup for the retry duration. 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}.Also when {@code retryAfterTimeUnit} is
+     * {@code null} and {@code retryAfterHeader} is not {@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 +141,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);
     }
 }