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,61 @@
 
 package com.azure.core.http.policy;
 
+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 pipeline policy that puts a UUID or user provided id in the request header. Azure uses the request id as
  * the unique identifier for the request.
  */
 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 = null;
+    }
+
+    /**
+     * 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}. {@code null} is
+     * valid value. It is suggested that this {@link Supplier} should provide unique value every time
+     * it is called. Example of these headers are 'x-ms-client-request-id', 'x-ms-correlation-request-id'.
+     */
+    public RequestIdPolicy(Supplier<HttpHeaders> requestIdSupplier) {
+        this.requestIdSupplier = requestIdSupplier;
+    }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+
+        if (Objects.nonNull(requestIdSupplier)) {
+            HttpHeaders httpHeaders = requestIdSupplier.get();
+            if (Objects.nonNull(httpHeaders) && httpHeaders.getSize() > 0) {
+                httpHeaders.stream().forEach(httpHeader -> {
+                    String requestIdHeaderValue = context.getHttpRequest().getHeaders().getValue(httpHeader.getName());
+                    if (requestIdHeaderValue == null) {
+                        context.getHttpRequest().getHeaders().put(httpHeader.getName(), httpHeader.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,88 @@
 
 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 static final String RETRY_AFTER_MS_HEADER = "retry-after-ms";
 
-    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
+     * Creates {@link RetryPolicy} with default {@link ExponentialBackoff} as {@link RetryStrategy} and use
      * 'retry-after-ms' in {@link HttpResponse} header for calculating retry delay.
      */
     public RetryPolicy() {
-        this(new RetryPolicyOptions(new ExponentialBackoff()));
+        this(new ExponentialBackoff(), RETRY_AFTER_MS_HEADER, ChronoUnit.MILLIS);
     }
 
     /**
-     * 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 use 'retry-after-ms' in
+     * {@link HttpResponse} header for calculating retry delay.
      *
-     * @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, RETRY_AFTER_MS_HEADER, ChronoUnit.MILLIS);
     }
 
     @Override
@@ -79,11 +108,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 +122,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/http/policy/RequestIdPolicyTests.java

@@ -108,7 +108,7 @@ public Mono<HttpResponse> send(HttpRequest request) {
                     return Mono.just(mockResponse);
                 }
             })
-            .policies(new RequestIdPolicy(), new RetryPolicy(new RetryPolicyOptions(new FixedDelay(1, Duration.of(0, ChronoUnit.SECONDS)))))
+            .policies(new RequestIdPolicy(), new RetryPolicy(new FixedDelay(1, Duration.of(0, ChronoUnit.SECONDS))))
             .build();
 
         pipeline.send(new HttpRequest(HttpMethod.GET, new URL("http://localhost/"))).block();