CosmosClient: Fixes missing Trace when converting HTTP Timeout to 503

Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs

@@ -49,7 +49,7 @@ await this.cosmosAuthorization.AddAuthorizationHeaderAsync(
 
             using (ITrace trace = Trace.GetRootTrace("Account Read", TraceComponent.Transport, TraceLevel.Info))
             {
-                IClientSideRequestStatistics stats = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace.Summary);
+                IClientSideRequestStatistics stats = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace);
 
                 try
                 {

Microsoft.Azure.Cosmos/src/Handler/TransportHandler.cs

@@ -93,7 +93,7 @@ internal async Task<ResponseMessage> ProcessMessageAsync(
 
             DocumentServiceRequest serviceRequest = request.ToDocumentServiceRequest();
 
-            ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, request.Trace.Summary);
+            ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, request.Trace);
             serviceRequest.RequestContext.ClientRequestStatistics = clientSideRequestStatisticsTraceDatum;
 
             //TODO: extrace auth into a separate handler

Microsoft.Azure.Cosmos/src/HttpClient/CosmosHttpClientCore.cs

@@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Resource.CosmosExceptions;
+    using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Cosmos.Tracing.TraceData;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Collections;
@@ -361,9 +362,11 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
                     }
                     catch (Exception e)
                     {
+                        ITrace trace = NoOpTrace.Singleton;
                         if (clientSideRequestStatistics is ClientSideRequestStatisticsTraceDatum datum)
                         {
                             datum.RecordHttpException(requestMessage, e, resourceType, requestStartTime);
+                            trace = datum.Trace;
                         }
                         bool isOutOfRetries = CosmosHttpClientCore.IsOutOfRetries(timeoutPolicy, startDateTimeUtc, timeoutEnumerator);
 
@@ -394,6 +397,7 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
                                                 ActivityId = System.Diagnostics.Trace.CorrelationManager.ActivityId.ToString(),
                                                 SubStatusCode = SubStatusCodes.TransportGenerated503
                                             },
+                                            trace: trace,
                                             innerException: e);
                                     }
 

Microsoft.Azure.Cosmos/src/Routing/ClientCollectionCache.cs

@@ -186,7 +186,7 @@ private async Task<ContainerProperties> ReadCollectionAsync(
                 {
                     headers.XDate = Rfc1123DateTimeCache.UtcNow();
 
-                    request.RequestContext.ClientRequestStatistics = clientSideRequestStatistics ?? new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace.Summary);
+                    request.RequestContext.ClientRequestStatistics = clientSideRequestStatistics ?? new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace);
                     if (clientSideRequestStatistics == null)
                     {
                         childTrace.AddDatum(

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs

@@ -316,7 +316,7 @@ private async Task<DocumentServiceResponse> ExecutePartitionKeyRangeReadChangeFe
                     }
 
                     request.Headers[HttpConstants.HttpHeaders.Authorization] = authorizationToken;
-                    request.RequestContext.ClientRequestStatistics = clientSideRequestStatistics ?? new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace.Summary);
+                    request.RequestContext.ClientRequestStatistics = clientSideRequestStatistics ?? new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace);
                     if (clientSideRequestStatistics == null)
                     {
                         childTrace.AddDatum("Client Side Request Stats", request.RequestContext.ClientRequestStatistics);

Microsoft.Azure.Cosmos/src/Tracing/TraceData/ClientSideRequestStatisticsTraceDatum.cs

@@ -33,9 +33,8 @@ internal sealed class ClientSideRequestStatisticsTraceDatum : TraceDatum, IClien
         private IReadOnlyList<StoreResponseStatistics> shallowCopyOfStoreResponseStatistics = null;
         private IReadOnlyList<HttpResponseStatistics> shallowCopyOfHttpResponseStatistics = null;
         private SystemUsageHistory systemUsageHistory = null;
-        public TraceSummary TraceSummary = null;
 
-        public ClientSideRequestStatisticsTraceDatum(DateTime startTime, TraceSummary summary)
+        public ClientSideRequestStatisticsTraceDatum(DateTime startTime, ITrace trace)
         {
             this.RequestStartTimeUtc = startTime;
             this.RequestEndTimeUtc = null;
@@ -45,7 +44,7 @@ public ClientSideRequestStatisticsTraceDatum(DateTime startTime, TraceSummary su
             this.FailedReplicas = new HashSet<TransportAddressUri>();
             this.RegionsContacted = new HashSet<(string, Uri)>();
             this.httpResponseStatistics = new List<HttpResponseStatistics>();
-            this.TraceSummary = summary;
+            this.Trace = trace;
         }
 
         public DateTime RequestStartTimeUtc { get; }
@@ -75,6 +74,10 @@ public IReadOnlyDictionary<string, AddressResolutionStatistics> EndpointToAddres
 
         public HashSet<(string, Uri)> RegionsContacted { get; }
 
+        public ITrace Trace { get; private set; }
+
+        public TraceSummary TraceSummary => this.Trace?.Summary;
+
         public IReadOnlyList<StoreResponseStatistics> StoreResponseStatisticsList
         {
             get
@@ -502,11 +505,11 @@ public HttpResponseStatistics(
                 if (responseMessage != null)
                 {
                     Headers headers = new Headers(GatewayStoreClient.ExtractResponseHeaders(responseMessage));
-                    this.ActivityId = headers.ActivityId ?? Trace.CorrelationManager.ActivityId.ToString();
+                    this.ActivityId = headers.ActivityId ?? System.Diagnostics.Trace.CorrelationManager.ActivityId.ToString();
                 }
                 else
                 {
-                    this.ActivityId = Trace.CorrelationManager.ActivityId.ToString();
+                    this.ActivityId = System.Diagnostics.Trace.CorrelationManager.ActivityId.ToString();
                 }
             }
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Benchmarks/RegionContactedInDiagnosticsBenchmark.cs

@@ -42,7 +42,7 @@ private ITrace CreateTestTraceTree()
 
         private TraceDatum GetDatumObject(string regionName1, string regionName2 = null)
         {
-            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary());
+            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, Trace.GetRootTrace(nameof(RegionContactedInDiagnosticsBenchmark)));
             Uri uri1 = new Uri("http://someUri1.com");
             datum.RegionsContacted.Add((regionName1, uri1));
             if (regionName2 != null)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/TraceWriterBaselineTests.TraceData.xml

@@ -187,7 +187,7 @@
         TraceForBaselineTesting rootTrace;
         using (rootTrace = TraceForBaselineTesting.GetRootTrace())
         {
-            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, new TraceSummary());
+            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, rootTrace);
 
             TransportAddressUri uri1 = new TransportAddressUri(new Uri("http://someUri1.com"));
             TransportAddressUri uri2 = new TransportAddressUri(new Uri("http://someUri2.com"));
@@ -405,7 +405,7 @@
         TraceForBaselineTesting rootTrace;
         using (rootTrace = TraceForBaselineTesting.GetRootTrace())
         {
-            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, new TraceSummary());
+            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, rootTrace);
             datum.ContactedReplicas.Add(default);
 
             TraceWriterBaselineTests.GetPrivateField<Dictionary<string, AddressResolutionStatistics>>(datum, "endpointToAddressResolutionStats").Add("asdf", default);
@@ -546,7 +546,7 @@
         TraceForBaselineTesting rootTrace;
         using (rootTrace = TraceForBaselineTesting.GetRootTrace())
         {
-            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, new TraceSummary());
+            ClientSideRequestStatisticsTraceDatum datum = new ClientSideRequestStatisticsTraceDatum(DateTime.MinValue, rootTrace);
             TraceWriterBaselineTests.SetEndRequestTime(datum,DateTime.MaxValue);
 
             HttpResponseStatistics httpResponseStatistics = new HttpResponseStatistics(