Show correct client request id on debug message (#235)

* Show correct client request id on debug message

* Support generated cmdlet

* Update code

* update code

* update code

* Update code after review

Author: dingmeng-xue

src/Common/CmdletInfoHandler.cs

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.WindowsAzure.Commands.Common.Extensions;
 using System;
 using System.Net.Http;
 using System.Threading;
@@ -64,11 +65,7 @@ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage reques
             }
             if (ClientRequestId != null)
             {
-                if (request.Headers.Contains("x-ms-client-request-id"))
-                {
-                    request.Headers.Remove("x-ms-client-request-id");
-                }
-                request.Headers.TryAddWithoutValidation("x-ms-client-request-id", ClientRequestId);
+                request.AddClientRequestId(ClientRequestId);
             }
             return base.SendAsync(request, cancellationToken);
         }

src/Common/Extensions/HttpRequestMessageExtensions.cs

@@ -0,0 +1,38 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System.Net.Http;
+
+namespace Microsoft.WindowsAzure.Commands.Common.Extensions
+{
+    public static class HttpRequestMessageExtensions
+    {
+        private const string HeaderNameClientRequestId = "x-ms-client-request-id";
+
+        /// <summary>
+        /// Add x-ms-client-request-id to headers of Http Request 
+        /// </summary>
+        /// <param name="request">Current Http Request</param>
+        /// <param name="clientRequestId">value of client request id</param>
+        public static void AddClientRequestId(this HttpRequestMessage request, string clientRequestId)
+        {
+            if (request.Headers.Contains(HeaderNameClientRequestId))
+            {
+                request.Headers.Remove(HeaderNameClientRequestId);
+            }
+            request.Headers.TryAddWithoutValidation(HeaderNameClientRequestId, clientRequestId);
+        }
+
+    }
+}

src/Common/RecordingTracingInterceptor.cs

@@ -14,6 +14,7 @@
 
 using Hyak.Common;
 using Microsoft.WindowsAzure.Commands.Common;
+using Microsoft.WindowsAzure.Commands.Common.Extensions;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using System;
 using System.Collections.Concurrent;
@@ -26,16 +27,19 @@ namespace Microsoft.Azure.ServiceManagement.Common.Models
     public class RecordingTracingInterceptor : Hyak.Common.ICloudTracingInterceptor
     {
 
-        public RecordingTracingInterceptor(ConcurrentQueue<string> queue, IList<Regex> matchers = null)
+        public RecordingTracingInterceptor(ConcurrentQueue<string> queue, IList<Regex> matchers = null, string clientRequestId = null)
         {
             MessageQueue = queue;
             Matchers = matchers;
+            this.clientRequestId = clientRequestId;
         }
 
         public ConcurrentQueue<string> MessageQueue { get; private set; }
 
         private IList<Regex> Matchers { get; set; }
 
+        private string clientRequestId;
+
         private void Write(string message, params object[] arguments)
         {
             if (arguments == null || arguments.Length == 0)
@@ -65,6 +69,12 @@ public void Enter(string invocationId, object instance, string method, IDictiona
 
         public void SendRequest(string invocationId, HttpRequestMessage request)
         {
+            // CmdletInfoHandler sets/updates x-ms-client-request-id during SendAsync() no matter if SDK sets x-ms-client-request-id.
+            // Update request here to ensure its value consistent with real result.
+            if (clientRequestId != null)
+            {
+                request.AddClientRequestId(clientRequestId);
+            }
             Write(GeneralUtilities.GetLog(request, Matchers));
         }
 

src/Common/Utilities/GeneralUtilities.cs

@@ -311,6 +311,11 @@ public static string GetHttpRequestLog(string method, string requestUri, HttpHea
 
         public static string GetLog(HttpResponseMessage response, IList<Regex> matchers = null)
         {
+            if (response == null)
+            {
+                return string.Empty;
+            }
+
             string body = response.Content == null ? string.Empty
                 : FormatString(response.Content.ReadAsStringAsync().Result);
 
@@ -323,6 +328,11 @@ public static string GetLog(HttpResponseMessage response, IList<Regex> matchers
 
         public static string GetLog(HttpRequestMessage request, IList<Regex> matchers = null)
         {
+            if (request == null)
+            {
+                return string.Empty;
+            }
+
             string body = request.Content == null ? string.Empty
                 : FormatString(request.Content.ReadAsStringAsync().Result);
 

src/ResourceManager/Version2016_09_01/AzureRMCmdlet.cs

@@ -324,7 +324,7 @@ protected override void SetupDebuggingTraces()
             ServiceClientTracing.IsEnabled = true;
             base.SetupDebuggingTraces();
             _serviceClientTracingInterceptor = _serviceClientTracingInterceptor
-                ?? new ServiceClientTracingInterceptor(DebugMessages, _matchers);
+                ?? new ServiceClientTracingInterceptor(DebugMessages, _matchers, _clientRequestId);
             ServiceClientTracing.AddTracingInterceptor(_serviceClientTracingInterceptor);
         }
 

src/ResourceManager/Version2016_09_01/ServiceClientTracingInterceptor.cs

@@ -14,6 +14,7 @@
 
 using Microsoft.Rest;
 using Microsoft.WindowsAzure.Commands.Common;
+using Microsoft.WindowsAzure.Commands.Common.Extensions;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
 using System;
 using System.Collections.Concurrent;
@@ -24,16 +25,19 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common
 {
     public class ServiceClientTracingInterceptor : IServiceClientTracingInterceptor
     {
-        public ServiceClientTracingInterceptor(ConcurrentQueue<string> queue, IList<Regex> matchers = null)
+        public ServiceClientTracingInterceptor(ConcurrentQueue<string> queue, IList<Regex> matchers = null, string clientRequestId = null)
         {
             MessageQueue = queue;
             Matchers = matchers;
+            this.clientRequestId = clientRequestId;
         }
 
         public ConcurrentQueue<string> MessageQueue { get; private set; }
 
         private IList<Regex> Matchers { get; set; }
 
+        private string clientRequestId;
+
         public void Configuration(string source, string name, string value)
         {
             // Ignore 
@@ -62,8 +66,13 @@ public void ReceiveResponse(string invocationId, System.Net.Http.HttpResponseMes
 
         public void SendRequest(string invocationId, System.Net.Http.HttpRequestMessage request)
         {
-            string requestAsString = request == null ? string.Empty : GeneralUtilities.GetLog(request, Matchers);
-            MessageQueue.CheckAndEnqueue(requestAsString);
+            // CmdletInfoHandler sets/updates x-ms-client-request-id during SendAsync() no matter if SDK sets x-ms-client-request-id.
+            // Update request here to ensure its value consistent with real result.
+            if (request != null && clientRequestId != null)
+            {
+                request.AddClientRequestId(clientRequestId);
+            }
+            MessageQueue.CheckAndEnqueue(GeneralUtilities.GetLog(request, Matchers));
         }
 
         public void TraceError(string invocationId, Exception exception)