Adding Options for Request Id Policy.

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

@@ -3,25 +3,64 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpHeaders;
 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;
 
 /**
- * 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.
+ * @see RequestIdPolicyOptions
  */
 public class RequestIdPolicy implements HttpPipelinePolicy {
     private static final String REQUEST_ID_HEADER = "x-ms-client-request-id";
 
+    private final RequestIdPolicyOptions requestIdPolicyOptions;
+
+    /**
+     * Creates {@link RequestIdPolicy}  with default behaviour.
+     */
+    public RequestIdPolicy() {
+        this.requestIdPolicyOptions = null;
+    }
+
+    /**
+     * Creates a {@link RequestIdPolicy} with the provided {@link RequestIdPolicyOptions}.
+     * User can specify {@link java.util.function.Function} to dynamically generated id for various id headers.
+     *
+     * @param requestIdPolicyOptions with given {@link RetryPolicyOptions}.
+     */
+    public RequestIdPolicy(RequestIdPolicyOptions requestIdPolicyOptions) {
+        this.requestIdPolicyOptions = requestIdPolicyOptions;
+    }
+
+
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
         String requestId = context.getHttpRequest().getHeaders().getValue(REQUEST_ID_HEADER);
-        if (requestId == null) {
-            context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());
+        if (!Objects.isNull(requestIdPolicyOptions)) {
+            HttpHeaders httpHeaders =  requestIdPolicyOptions.getIdHeaderSupplier().get();
+            if (!Objects.isNull(httpHeaders)) {
+                httpHeaders.stream().forEach(httpHeader -> {
+                    String requestIdHeaderValue = context.getHttpRequest().getHeaders().getValue(httpHeader.getName());
+                    if (requestIdHeaderValue == null) {
+                        context.getHttpRequest().getHeaders().put(httpHeader.getName(), httpHeader.getValue());
+                    }
+            });
+            }else {
+                if (requestId == null) {
+                    context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());
+                }
+            }
+        } else {
+            if (requestId == null) {
+                context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());
+            }
         }
         return next.process();
     }

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

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http.policy;
+
+import com.azure.core.annotation.Immutable;
+import com.azure.core.http.HttpHeaders;
+
+import java.util.Objects;
+import java.util.function.Supplier;
+
+/**
+ * Immutable Configuration options for {@link RequestIdPolicy}.
+ * User can provide a {@link Supplier} to generated id for various id headers to be used in
+ * {@link com.azure.core.http.HttpRequest}. Examples of such headers are 'x-ms-client-request-id
+ * or 'x-ms-correlation-request-id' etc.
+ */
+@Immutable
+public class RequestIdPolicyOptions {
+
+    private final Supplier<HttpHeaders> idHeaderSupplier;
+
+    /**
+     * Creates {@link RequestIdPolicyOptions} with provided {@code idHeaderSupplier }.
+     *
+     * @param idHeaderSupplier {@link Supplier} to provide dynamically generated id for various request id headers.
+     * Examples of such headers are 'x-ms-client-request-id or 'x-ms-correlation-request-id' etc.
+     */
+    public RequestIdPolicyOptions(Supplier<HttpHeaders> idHeaderSupplier) {
+        this.idHeaderSupplier = Objects.requireNonNull(idHeaderSupplier,
+            "'idHeaderSupplier' cannot be null.");
+    }
+
+    /**
+     * Get the provided {@link Supplier} which dynamically generate id for various request id headers.
+     *
+     * @return {@link Supplier} to provide dynamically generated id for various request id headers
+     */
+    public Supplier<HttpHeaders> getIdHeaderSupplier() {
+        return idHeaderSupplier;
+    }
+}