CosmosClient: Fixes missing Trace when converting HTTP Timeout to 503

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -6473,8 +6473,10 @@ await this.cosmosAuthorization.AddAuthorizationHeaderAsync(
                     return request;
                 }
 
-                AccountProperties databaseAccount = await gatewayModel.GetDatabaseAccountAsync(CreateRequestMessage,
-                                                                                               clientSideRequestStatistics: null);
+                AccountProperties databaseAccount = await gatewayModel.GetDatabaseAccountAsync(
+                    CreateRequestMessage,
+                    clientSideRequestStatistics: null,
+                    trace: null);
                 this.UseMultipleWriteLocations = this.ConnectionPolicy.UseMultipleWriteLocations && databaseAccount.EnableMultipleWriteLocations;
 
                 if (this.queryPartitionProvider.IsValueCreated)

Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs

@@ -59,6 +59,7 @@ await this.cosmosAuthorization.AddAuthorizationHeaderAsync(
                         resourceType: ResourceType.DatabaseAccount,
                         timeoutPolicy: HttpTimeoutPolicyControlPlaneRead.Instance,
                         clientSideRequestStatistics: stats,
+                        trace: trace,
                         cancellationToken: default))
                     using (DocumentServiceResponse documentServiceResponse = await ClientExtensions.ParseResponseAsync(responseMessage))
                     {

Microsoft.Azure.Cosmos/src/GatewayStoreClient.cs

@@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Handlers;
+    using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Cosmos.Tracing.TraceData;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Collections;
@@ -78,13 +79,15 @@ internal Task<HttpResponseMessage> SendHttpAsync(
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken = default)
         {
             return this.httpClient.SendHttpAsync(
                 createRequestMessageAsync: requestMessage,
                 resourceType: resourceType,
                 timeoutPolicy: timeoutPolicy,
                 clientSideRequestStatistics: clientSideRequestStatistics,
+                trace: trace,
                 cancellationToken: cancellationToken);
         }
 
@@ -354,6 +357,7 @@ private Task<HttpResponseMessage> InvokeClientAsync(
                 resourceType,
                 HttpTimeoutPolicy.GetTimeoutPolicy(request),
                 request.RequestContext.ClientRequestStatistics,
+                trace: null,
                 cancellationToken);
         }
     }

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

@@ -93,9 +93,11 @@ await GatewayStoreModel.ApplySessionTokenAsync(
             return response;
         }
 
-        public virtual async Task<AccountProperties> GetDatabaseAccountAsync(Func<ValueTask<HttpRequestMessage>> requestMessage,
-                                                        IClientSideRequestStatistics clientSideRequestStatistics,
-                                                        CancellationToken cancellationToken = default)
+        public virtual async Task<AccountProperties> GetDatabaseAccountAsync(
+            Func<ValueTask<HttpRequestMessage>> requestMessage,
+            IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
+            CancellationToken cancellationToken = default)
         {
             AccountProperties databaseAccount = null;
 
@@ -105,6 +107,7 @@ public virtual async Task<AccountProperties> GetDatabaseAccountAsync(Func<ValueT
                 ResourceType.DatabaseAccount,
                 HttpTimeoutPolicyControlPlaneRead.Instance,
                 clientSideRequestStatistics,
+                trace,
                 cancellationToken))
             {
                 using (DocumentServiceResponse documentServiceResponse = await ClientExtensions.ParseResponseAsync(responseMessage))

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

@@ -23,13 +23,15 @@ public abstract Task<HttpResponseMessage> GetAsync(
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken);
 
         public abstract Task<HttpResponseMessage> SendHttpAsync(
             Func<ValueTask<HttpRequestMessage>> createRequestMessageAsync,
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken);
 
         protected abstract void Dispose(bool disposing);

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;
@@ -261,6 +262,7 @@ public override Task<HttpResponseMessage> GetAsync(
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken)
         {
             if (uri == null)
@@ -292,6 +294,7 @@ ValueTask<HttpRequestMessage> CreateRequestMessage()
                 resourceType,
                 timeoutPolicy,
                 clientSideRequestStatistics,
+                trace,
                 cancellationToken);
         }
 
@@ -300,6 +303,7 @@ public override Task<HttpResponseMessage> SendHttpAsync(
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken)
         {
             if (createRequestMessageAsync == null)
@@ -312,6 +316,7 @@ public override Task<HttpResponseMessage> SendHttpAsync(
                 resourceType,
                 timeoutPolicy,
                 clientSideRequestStatistics,
+                trace,
                 cancellationToken);
         }
 
@@ -320,6 +325,7 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
             ResourceType resourceType,
             HttpTimeoutPolicy timeoutPolicy,
             IClientSideRequestStatistics clientSideRequestStatistics,
+            ITrace trace,
             CancellationToken cancellationToken)
         {
             DateTime startDateTimeUtc = DateTime.UtcNow;
@@ -394,6 +400,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/GatewayAddressCache.cs

@@ -613,6 +613,7 @@ private async Task<DocumentServiceResponse> GetMasterAddressesViaGatewayAsync(
                     resourceType: resourceType,
                     timeoutPolicy: HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,
                     clientSideRequestStatistics: request.RequestContext?.ClientRequestStatistics,
+                    trace: trace,
                     cancellationToken: default))
                 {
                     DocumentServiceResponse documentServiceResponse = await ClientExtensions.ParseResponseAsync(httpResponseMessage);
@@ -694,6 +695,7 @@ private async Task<DocumentServiceResponse> GetServerAddressesViaGatewayAsync(
                     resourceType: ResourceType.Document,
                     timeoutPolicy: HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,
                     clientSideRequestStatistics: request.RequestContext?.ClientRequestStatistics,
+                    trace: trace,
                     cancellationToken: default))
                 {
                     DocumentServiceResponse documentServiceResponse = await ClientExtensions.ParseResponseAsync(httpResponseMessage);

Microsoft.Azure.Cosmos/src/Telemetry/ClientTelemetryProcessor.cs

@@ -114,6 +114,7 @@ await this.tokenProvider.AddAuthorizationHeaderAsync(
                                                     ResourceType.Telemetry,
                                                     HttpTimeoutPolicyNoRetry.Instance,
                                                     null,
+                                                    null,
                                                     cancellationToken);
 
                 if (!response.IsSuccessStatusCode)

Microsoft.Azure.Cosmos/src/Telemetry/VmMetadataApiHandler.cs

@@ -81,6 +81,7 @@ static ValueTask<HttpRequestMessage> CreateRequestMessage()
                                     resourceType: ResourceType.Telemetry,
                                     timeoutPolicy: HttpTimeoutPolicyNoRetry.Instance,
                                     clientSideRequestStatistics: null,
+                                    trace: null,
                                     cancellationToken: default);
 
                 azMetadata = await VmMetadataApiHandler.ProcessResponseAsync(response);