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/CosmosHttpClientCoreTests.cs

@@ -42,14 +42,17 @@ static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellati
 
             HttpRequestMessage httpRequestMessage = new HttpRequestMessage(HttpMethod.Post, new Uri("http://localhost"));
 
-            HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
-                new ValueTask<HttpRequestMessage>(httpRequestMessage),
-                ResourceType.Collection,
-                timeoutPolicy: HttpTimeoutPolicyDefault.Instance,
-                new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow,new TraceSummary()),
-                default);
-
-            Assert.AreEqual(httpRequestMessage, responseMessage.RequestMessage);
+            using (ITrace trace = Trace.GetRootTrace(nameof(ResponseMessageHasRequestMessageAsync)))
+            {
+                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+                    new ValueTask<HttpRequestMessage>(httpRequestMessage),
+                    ResourceType.Collection,
+                    timeoutPolicy: HttpTimeoutPolicyDefault.Instance,
+                    new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
+                    default);
+
+                Assert.AreEqual(httpRequestMessage, responseMessage.RequestMessage);
+            }
         }
 
         [TestMethod]
@@ -111,15 +114,18 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
                 HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
                 using CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
 
-                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+                using (ITrace trace = Trace.GetRootTrace(nameof(RetryTransientIssuesTestAsync)))
+                {
+                    HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
                     new ValueTask<HttpRequestMessage>(
                         result: new HttpRequestMessage(HttpMethod.Get, new Uri("http://localhost"))),
                         resourceType: ResourceType.Collection,
                         timeoutPolicy: currentTimeoutPolicy.Key,
-                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                         cancellationToken: default);
 
-                Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+                    Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+                }
             }
         }
 
@@ -178,15 +184,18 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
             HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
             using CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
 
-            HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+            using (ITrace trace = Trace.GetRootTrace(nameof(RetryTransient408sTestAsync)))
+            {
+                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
                 new ValueTask<HttpRequestMessage>(
                     result: new HttpRequestMessage(HttpMethod.Get, new Uri("http://localhost"))),
                     resourceType: ResourceType.Collection,
                     timeoutPolicy: HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,
-                    clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                    clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                     cancellationToken: default);
 
-            Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+                Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+            }
         }
 
         [TestMethod]
@@ -216,15 +225,18 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
             HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
             using CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
 
-            HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+            using (ITrace trace = Trace.GetRootTrace(nameof(DoesNotRetryTransient408sOnDefaultPolicyTestAsync)))
+            {
+                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
                 new ValueTask<HttpRequestMessage>(
                     result: new HttpRequestMessage(HttpMethod.Get, new Uri("http://localhost"))),
                     resourceType: ResourceType.Collection,
                     timeoutPolicy: HttpTimeoutPolicyControlPlaneRead.Instance,
-                    clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                    clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                     cancellationToken: default);
-
-            Assert.AreEqual(HttpStatusCode.RequestTimeout, responseMessage.StatusCode, "Should be a request timeout");
+
+                Assert.AreEqual(HttpStatusCode.RequestTimeout, responseMessage.StatusCode, "Should be a request timeout");
+            }
         }
 
         [TestMethod]
@@ -245,13 +257,16 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
 
             try
             {
-                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+                using (ITrace trace = Trace.GetRootTrace(nameof(Retry3TimesOnDefaultPolicyTestAsync)))
+                {
+                    HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
                     new ValueTask<HttpRequestMessage>(
                         result: new HttpRequestMessage(HttpMethod.Get, new Uri("http://localhost"))),
                         resourceType: ResourceType.Collection,
                         timeoutPolicy: HttpTimeoutPolicyDefault.Instance,
-                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                         cancellationToken: default);
+                }
             }
             catch (Exception ex)
             {
@@ -282,14 +297,16 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
 
                 try
                 {
-
-                    HttpResponseMessage responseMessage1 = await cosmoshttpClient.SendHttpAsync(() =>
+                    using (ITrace trace = Trace.GetRootTrace(nameof(NoRetryOnNoRetryPolicyTestAsync)))
+                    {
+                        HttpResponseMessage responseMessage1 = await cosmoshttpClient.SendHttpAsync(() =>
                         new ValueTask<HttpRequestMessage>(
                             result: new HttpRequestMessage(method, new Uri("http://localhost"))),
                             resourceType: resourceType,
                             timeoutPolicy: timeoutPolicy,
-                            clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                            clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                             cancellationToken: default);
+                    }
                 }
                 catch (Exception e)
                 {
@@ -301,6 +318,9 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
                         CosmosException cosmosException = (CosmosException)e;
                         Assert.AreEqual(cosmosException.StatusCode, System.Net.HttpStatusCode.ServiceUnavailable);
                         Assert.AreEqual((int)cosmosException.SubStatusCode,(int)SubStatusCodes.TransportGenerated503);
+
+                        Assert.IsNotNull(cosmosException.Trace);
+                        Assert.AreNotEqual(cosmosException.Trace, NoOpTrace.Singleton);
                     }
                 }
 
@@ -344,13 +364,16 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
 
             try
             {
-                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+                using (ITrace trace = Trace.GetRootTrace(nameof(NoRetryOnNoRetryPolicyTestAsync)))
+                {
+                    HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
                     new ValueTask<HttpRequestMessage>(
                         result: new HttpRequestMessage(HttpMethod.Get, new Uri("http://localhost"))),
                         resourceType: ResourceType.Collection,
                         timeoutPolicy: HttpTimeoutPolicyNoRetry.Instance,
-                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
+                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
                         cancellationToken: default);
+                }
             }
             catch (Exception ex)
             {
@@ -401,15 +424,18 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
             HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
             using CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
 
-            HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
-                new ValueTask<HttpRequestMessage>(
-                    result: new HttpRequestMessage(HttpMethod.Post, new Uri("http://localhost"))),
-                    resourceType: ResourceType.Document,
-                    timeoutPolicy: HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,
-                    clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
-                    cancellationToken: default);
+            using (ITrace trace = Trace.GetRootTrace(nameof(RetryTransientIssuesForQueryPlanTestAsync)))
+            {
+                HttpResponseMessage responseMessage = await cosmoshttpClient.SendHttpAsync(() =>
+                    new ValueTask<HttpRequestMessage>(
+                        result: new HttpRequestMessage(HttpMethod.Post, new Uri("http://localhost"))),
+                        resourceType: ResourceType.Document,
+                        timeoutPolicy: HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,
+                        clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace),
+                        cancellationToken: default);
 
-            Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+                Assert.AreEqual(HttpStatusCode.OK, responseMessage.StatusCode);
+            }
         }
 
         [TestMethod]

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayAccountReaderTests.cs

@@ -57,15 +57,20 @@ public async Task DocumentClient_BuildHttpClientFactory_WithHandler()
                 receivedResponseEventArgs: null);
 
             Assert.IsNotNull(httpClient);
-            HttpResponseMessage response = await httpClient.GetAsync(
-                uri: new Uri("https://localhost"),
-                additionalHeaders: new RequestNameValueCollection(),
-                resourceType: ResourceType.Document,
-                timeoutPolicy: HttpTimeoutPolicyDefault.InstanceShouldThrow503OnTimeout,
-                clientSideRequestStatistics: new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, new TraceSummary()),
-                cancellationToken: default);
-
-            Assert.AreEqual(HttpStatusCode.Conflict, response.StatusCode);
+
+            using (ITrace trace = Trace.GetRootTrace(nameof(DocumentClient_BuildHttpClientFactory_WithHandler)))
+            {
+                IClientSideRequestStatistics stats = new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace);
+                HttpResponseMessage response = await httpClient.GetAsync(
+                    uri: new Uri("https://localhost"),
+                    additionalHeaders: new RequestNameValueCollection(),
+                    resourceType: ResourceType.Document,
+                    timeoutPolicy: HttpTimeoutPolicyDefault.InstanceShouldThrow503OnTimeout,
+                    clientSideRequestStatistics: stats,
+                    cancellationToken: default);
+
+                Assert.AreEqual(HttpStatusCode.Conflict, response.StatusCode);
+            }
         }
 
         [TestMethod]