Customizing azure-core Http Header in RetryPolicy #6217

sdk/appconfiguration/azure-data-appconfiguration/src/main/java/com/azure/data/appconfiguration/ConfigurationClientBuilder.java

@@ -11,20 +11,18 @@
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.http.policy.RequestIdPolicy;
 import com.azure.core.http.policy.RetryPolicy;
-import com.azure.core.http.policy.RetryPolicyOptions;
-import com.azure.core.http.policy.ExponentialBackoff;
 import com.azure.core.http.policy.AddDatePolicy;
 import com.azure.core.http.policy.UserAgentPolicy;
 import com.azure.core.http.policy.HttpPolicyProviders;
 import com.azure.core.http.policy.HttpLogOptions;
 import com.azure.core.util.logging.ClientLogger;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpClient;
 import com.azure.data.appconfiguration.implementation.ConfigurationClientCredentials;
 import com.azure.data.appconfiguration.implementation.ConfigurationCredentialsPolicy;
 import com.azure.data.appconfiguration.models.ConfigurationSetting;
 import com.azure.core.util.Configuration;
-import com.azure.core.http.HttpClient;
-import com.azure.core.http.HttpHeaders;
-import com.azure.core.http.HttpPipeline;
 import com.azure.core.util.CoreUtils;
 
 import java.net.MalformedURLException;
@@ -79,9 +77,7 @@ public final class ConfigurationClientBuilder {
     private static final String APP_CONFIG_PROPERTIES = "azure-appconfig.properties";
     private static final String NAME = "name";
     private static final String VERSION = "version";
-    private static final String RETRY_AFTER_MS_HEADER = "retry-after-ms";
-    private static final RetryPolicy DEFAULT_RETRY_POLICY = new RetryPolicy(
-        new RetryPolicyOptions(new ExponentialBackoff(), RETRY_AFTER_MS_HEADER, ChronoUnit.MILLIS));
+    private static final RetryPolicy DEFAULT_RETRY_POLICY = new RetryPolicy("retry-after-ms", ChronoUnit.MILLIS);
 
     private final ClientLogger logger = new ClientLogger(ConfigurationClientBuilder.class);
     private final List<HttpPipelinePolicy> policies;

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

@@ -3,22 +3,63 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpHeader;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpResponse;
 import reactor.core.publisher.Mono;
 
+import java.util.Objects;
 import java.util.UUID;
+import java.util.function.Supplier;
 
 /**
- * The pipeline policy that puts a UUID in the request header. Azure uses the request id as
- * the unique identifier for the request.
+ * The pipeline policy that adds request id header in {@link HttpRequest} once. These id does not change
+ * when request is retried. Azure uses the request id as the unique identifier for the request.
+ * Example of these headers are 'x-ms-client-request-id' and 'x-ms-correlation-request-id'.
  */
 public class RequestIdPolicy implements HttpPipelinePolicy {
     private static final String REQUEST_ID_HEADER = "x-ms-client-request-id";
+    private final Supplier<HttpHeaders> requestIdSupplier;
+
+    /**
+     * Creates default {@link RequestIdPolicy}.
+     */
+    public RequestIdPolicy() {
+        requestIdSupplier = () -> new HttpHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());
+    }
+
+    /**
+     * Creates {@link RequestIdPolicy} with provided {@link Supplier} to dynamically generate request id for each
+     * {@link HttpRequest}.
+     *
+     * @param requestIdSupplier to dynamically generate to request id for each {@link HttpRequest}. It is suggested
+     * that this {@link Supplier} provides unique value every time it is called.
+     * Example of these headers are 'x-ms-client-request-id', 'x-ms-correlation-request-id'.
+     *
+     * @throws NullPointerException when {@code requestIdSupplier} is {@code null}.
+     */
+    public RequestIdPolicy(Supplier<HttpHeaders> requestIdSupplier) {
+        this.requestIdSupplier = Objects.requireNonNull(requestIdSupplier, "'requestIdSupplier' must not be null");
+    }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+
+        HttpHeaders httpHeaders = requestIdSupplier.get();
+        if (Objects.nonNull(httpHeaders) && httpHeaders.getSize() > 0) {
+            for (HttpHeader header : httpHeaders) {
+                String requestIdHeaderValue = context.getHttpRequest().getHeaders().getValue(header.getName());
+                if (requestIdHeaderValue == null) {
+                    context.getHttpRequest().getHeaders().put(header.getName(), header.getValue());
+                }
+            }
+            return next.process();
+        }
+
+        // If we were not able to set client provided Request ID header, we will set default 'REQUEST_ID_HEADER'.
         String requestId = context.getHttpRequest().getHeaders().getValue(REQUEST_ID_HEADER);
         if (requestId == null) {
             context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());

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

@@ -3,59 +3,89 @@
 
 package com.azure.core.http.policy;
 
-import static com.azure.core.util.CoreUtils.isNullOrEmpty;
-
 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.time.temporal.ChronoUnit;
 import java.util.Objects;
 import com.azure.core.util.logging.ClientLogger;
 import reactor.core.publisher.Mono;
 
+import static com.azure.core.util.CoreUtils.isNullOrEmpty;
+
 import java.time.Duration;
 
 /**
  * A pipeline policy that retries when a recoverable HTTP error occurs.
- * @see RetryPolicyOptions
  */
 public class RetryPolicy implements HttpPipelinePolicy {
 
     private final ClientLogger logger = new ClientLogger(RetryPolicy.class);
 
-    private final RetryPolicyOptions retryPolicyOptions;
+    private final RetryStrategy retryStrategy;
+    private final String retryAfterHeader;
+    private final ChronoUnit retryAfterTimeUnit;
 
     /**
-     * Creates {@link RetryPolicy} with default {@link ExponentialBackoff} as {@link RetryStrategy}and use
-     * 'retry-after-ms' in {@link HttpResponse} header for calculating retry delay.
+     * Creates {@link RetryPolicy} with default {@link ExponentialBackoff} as {@link RetryStrategy} and ignore the
+     * delay provided in response header.
      */
     public RetryPolicy() {
-        this(new RetryPolicyOptions(new ExponentialBackoff()));
+        this(new ExponentialBackoff(), null, null);
     }
 
     /**
-     * Creates a RetryPolicy with the provided {@link RetryStrategy}.
+     * Creates {@link RetryPolicy} with default {@link ExponentialBackoff} as {@link RetryStrategy} and use
+     * provided {@code retryAfterHeader} in {@link HttpResponse} headers for calculating retry delay.
+     *
+     * @param retryAfterHeader The HTTP header, such as 'Retry-After' or 'x-ms-retry-after-ms', to lookup for the
+     * retry delay. If the value is {@code null}, {@link RetryPolicy} will use the retry strategy to compute the delay
+     * and ignore the delay provided in response header.
+     * @param  retryAfterTimeUnit The time unit to use when applying the retry delay. {@code null} is valid if, and only
+     * if, {@code retryAfterHeader} is {@code null}.
+     * @throws NullPointerException When {@code retryAfterTimeUnit} is {@code null} and {@code retryAfterHeader} is
+     * not {@code null}.
+     */
+    public RetryPolicy(String retryAfterHeader, ChronoUnit retryAfterTimeUnit) {
+        this(new ExponentialBackoff(), retryAfterHeader, retryAfterTimeUnit);
+    }
+
+    /**
+     * Creates {@link RetryPolicy} with the provided {@link RetryStrategy} and default {@link ExponentialBackoff}
+     * as {@link RetryStrategy}. It will use provided {@code retryAfterHeader} in {@link HttpResponse} headers for
+     * calculating retry delay.
      *
      * @param retryStrategy The {@link RetryStrategy} used for retries.
+     * @param retryAfterHeader The HTTP header, such as 'Retry-After' or 'x-ms-retry-after-ms', to lookup for the
+     * retry delay. If the value is {@code null}, {@link RetryPolicy} will use the retry strategy to compute the delay
+     * and ignore the delay provided in response header.
+     * @param  retryAfterTimeUnit The time unit to use when applying the retry delay. {@code null} is valid if, and only
+     * if, {@code retryAfterHeader} is {@code null}.
+     *
+     * @throws NullPointerException When {@code retryStrategy} is {@code null}. Also when {@code retryAfterTimeUnit}
+     * is {@code null} and {@code retryAfterHeader} is not {@code null}.
      */
-    public RetryPolicy(RetryStrategy retryStrategy) {
-        Objects.requireNonNull(retryStrategy, "'retryStrategy' cannot be null");
-        this.retryPolicyOptions = new RetryPolicyOptions(retryStrategy);
+    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.");
+        }
     }
 
     /**
-     * Creates a {@link RetryPolicy} with the provided {@link RetryPolicyOptions}.
+     * Creates a {@link RetryPolicy} with the provided {@link RetryStrategy} and ignore the delay provided in
+     * response header.
      *
-     * @param retryPolicyOptions with given {@link RetryPolicyOptions}.
-     * @throws NullPointerException if {@code retryPolicyOptions} or  {@code retryPolicyOptions getRetryStrategy }
-     * is {@code null}.
+     * @param retryStrategy The {@link RetryStrategy} used for retries.
+     *
+     * @throws NullPointerException When {@code retryStrategy} is {@code null}.
      */
-    public RetryPolicy(RetryPolicyOptions retryPolicyOptions) {
-        this.retryPolicyOptions = Objects.requireNonNull(retryPolicyOptions,
-            "'retryPolicyOptions' cannot be null.");
-        Objects.requireNonNull(retryPolicyOptions.getRetryStrategy(),
-            "'retryPolicyOptions.retryStrategy' cannot be null.");
+    public RetryPolicy(RetryStrategy retryStrategy) {
+        this(retryStrategy, null, null);
     }
 
     @Override
@@ -79,11 +109,11 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, f
                 }
             })
             .onErrorResume(err -> {
-                int maxRetries = retryPolicyOptions.getRetryStrategy().getMaxRetries();
+                int maxRetries = retryStrategy.getMaxRetries();
                 if (tryCount < maxRetries) {
                     logger.verbose("[Error Resume] Try count: {}, Error: {}", tryCount, err);
                     return attemptAsync(context, next, originalHttpRequest, tryCount + 1)
-                        .delaySubscription(retryPolicyOptions.getRetryStrategy().calculateRetryDelay(tryCount));
+                        .delaySubscription(retryStrategy.calculateRetryDelay(tryCount));
                 } else {
                     return Mono.error(new RuntimeException(
                         String.format("Max retries %d times exceeded. Error Details: %s", maxRetries, err.getMessage()),
@@ -93,37 +123,37 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, f
     }
 
     private boolean shouldRetry(HttpResponse response, int tryCount) {
-        return tryCount < retryPolicyOptions.getRetryStrategy().getMaxRetries()
-            && retryPolicyOptions.getRetryStrategy().shouldRetry(response);
+        return tryCount < retryStrategy.getMaxRetries()
+            && retryStrategy.shouldRetry(response);
     }
 
     /**
      * Determines the delay duration that should be waited before retrying.
      * @param response HTTP response
      * @return If the HTTP response has a retry-after-ms header that will be returned,
-     *     otherwise the duration used during the construction of the policy.
+     * otherwise the duration used during the construction of the policy.
      */
     private Duration determineDelayDuration(HttpResponse response, int tryCount) {
         int code = response.getStatusCode();
 
         // Response will not have a retry-after-ms header.
         if (code != 429        // too many requests
             && code != 503) {  // service unavailable
-            return retryPolicyOptions.getRetryStrategy().calculateRetryDelay(tryCount);
+            return retryStrategy.calculateRetryDelay(tryCount);
         }
 
         String retryHeaderValue = null;
 
-        if (!isNullOrEmpty(retryPolicyOptions.getRetryAfterHeader())) {
-            retryHeaderValue = response.getHeaderValue(retryPolicyOptions.getRetryAfterHeader());
+        if (!isNullOrEmpty(this.retryAfterHeader)) {
+            retryHeaderValue = response.getHeaderValue(this.retryAfterHeader);
         }
 
         // Retry header is missing or empty, return the default delay duration.
         if (isNullOrEmpty(retryHeaderValue)) {
-            return retryPolicyOptions.getRetryStrategy().calculateRetryDelay(tryCount);
+            return this.retryStrategy.calculateRetryDelay(tryCount);
         }
 
         // Use the response delay duration, the server returned it for a reason.
-        return Duration.of(Integer.parseInt(retryHeaderValue), retryPolicyOptions.getRetryAfterTimeUnit());
+        return Duration.of(Integer.parseInt(retryHeaderValue), this.retryAfterTimeUnit);
     }
 }

sdk/core/azure-core/src/test/java/com/azure/core/UserAgentTests.java

@@ -12,7 +12,6 @@
 import com.azure.core.http.MockHttpResponse;
 import com.azure.core.http.policy.FixedDelay;
 import com.azure.core.http.policy.RetryPolicy;
-import com.azure.core.http.policy.RetryPolicyOptions;
 import com.azure.core.http.policy.UserAgentPolicy;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.Context;
@@ -58,7 +57,7 @@ public void userAgentPolicyAfterRetryPolicy(UserAgentPolicy userAgentPolicy, Str
         HttpPipeline pipeline = new HttpPipelineBuilder()
             .httpClient(new RetryValidationHttpClient(request ->
                 assertEquals(expected, request.getHeaders().getValue(USER_AGENT))))
-            .policies(new RetryPolicy(new RetryPolicyOptions(new FixedDelay(5, Duration.ofMillis(10)))))
+            .policies(new RetryPolicy(new FixedDelay(5, Duration.ofMillis(10))))
             .policies(userAgentPolicy)
             .build();