Internal ITrace: Adds wiring through EnsureValidClientAsync() (#2307)

This wires ITrace through the EnsureValidClientAsync. This is needed to see if the request latency is caused by the SDK initialization logic and warming up of the caches. Co-authored-by: j82w <j82w@users.noreply.github.com>
Azure · Mar 18, 2021 · 79b257e · 79b257e
1 parent 8ad6cd3
commit 79b257e
Show file tree

Hide file tree

Showing 31 changed files with 1,151 additions and 138 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/BulkExecutionRetryPolicy.cs b/Microsoft.Azure.Cosmos/src/BulkExecutionRetryPolicy.cs
@@ -106,7 +106,7 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
                     || subStatusCode == SubStatusCodes.CompletingSplit
                     || subStatusCode == SubStatusCodes.CompletingPartitionMigration)
                 {
-                    PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+                    PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
                     string containerRid = await this.container.GetCachedRIDAsync(
                         forceRefresh: false, 
                         NoOpTrace.Singleton, 

diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeed/StandByFeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/ChangeFeed/StandByFeedIteratorCore.cs
@@ -114,7 +114,7 @@ internal async Task<Tuple<string, ResponseMessage>> ReadNextInternalAsync(ITrace
 
             if (this.compositeContinuationToken == null)
             {
-                PartitionKeyRangeCache pkRangeCache = await this.clientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+                PartitionKeyRangeCache pkRangeCache = await this.clientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(trace);
                 this.containerRid = await this.container.GetCachedRIDAsync(
                     forceRefresh: false, 
                     trace, 

diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedProcessorCore.cs b/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedProcessorCore.cs
@@ -81,7 +81,7 @@ private async Task InitializeAsync()
                 default);
             string monitoredDatabaseAndContainerRid = await this.monitoredContainer.GetMonitoredDatabaseAndContainerRidAsync();
             string leaseContainerPrefix = this.monitoredContainer.GetLeasePrefix(this.changeFeedLeaseOptions.LeasePrefix, monitoredDatabaseAndContainerRid);
-            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.monitoredContainer.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.monitoredContainer.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             if (this.documentServiceLeaseStoreManager == null)
             {
                 this.documentServiceLeaseStoreManager = await DocumentServiceLeaseStoreManagerBuilder.InitializeAsync(this.monitoredContainer, this.leaseContainer, leaseContainerPrefix, this.instanceName).ConfigureAwait(false);

diff --git a/...Azure.Cosmos/src/ChangeFeedProcessor/LeaseManagement/DocumentServiceLeaseManagerCosmos.cs b/...Azure.Cosmos/src/ChangeFeedProcessor/LeaseManagement/DocumentServiceLeaseManagerCosmos.cs
@@ -65,7 +65,7 @@ public override async Task<DocumentServiceLease> AcquireAsync(DocumentServiceLea
                         throw tryInitializeContainerRId.Exception.InnerException;
                     }
 
-                    this.partitionKeyRangeCache = await this.monitoredContainer.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+                    this.partitionKeyRangeCache = await this.monitoredContainer.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
                 }
 
                 PartitionKeyRange partitionKeyRange = await this.partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(

diff --git a/Microsoft.Azure.Cosmos/src/DocumentClient.cs b/Microsoft.Azure.Cosmos/src/DocumentClient.cs
diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs b/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs
@@ -185,7 +185,7 @@ public override Documents.ShouldRetryResult HandleChangeFeedNotModified(Response
                 return FeedRangeCompositeContinuation.NoRetry;
             }
 
-            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await containerCore.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await containerCore.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             IReadOnlyList<Documents.PartitionKeyRange> resolvedRanges = await this.TryGetOverlappingRangesAsync(partitionKeyRangeCache, this.CurrentToken.Range.Min, this.CurrentToken.Range.Max, forceRefresh: true);
             if (resolvedRanges.Count > 0)
             {

diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRangeExtractor.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRangeExtractor.cs
@@ -23,7 +23,7 @@ public FeedRangePartitionKeyRangeExtractor(ContainerInternal container)
 
         public async Task<IReadOnlyList<Documents.Routing.Range<string>>> VisitAsync(FeedRangePartitionKey feedRange, CancellationToken cancellationToken = default)
         {
-            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             PartitionKeyDefinition partitionKeyDefinition = await this.container.GetPartitionKeyDefinitionAsync(cancellationToken);
             return await feedRange.GetEffectiveRangesAsync(
                 partitionKeyRangeCache,
@@ -38,7 +38,7 @@ await this.container.GetCachedRIDAsync(
         public async Task<IReadOnlyList<Documents.Routing.Range<string>>> VisitAsync(FeedRangePartitionKeyRange feedRange, CancellationToken cancellationToken = default)
         {
             // Migration from PKRangeId scenario
-            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             return await feedRange.GetEffectiveRangesAsync(
                 routingMapProvider: partitionKeyRangeCache,
                 containerRid: await this.container.GetCachedRIDAsync(
@@ -51,7 +51,7 @@ await this.container.GetCachedRIDAsync(
 
         public async Task<IReadOnlyList<Documents.Routing.Range<string>>> VisitAsync(FeedRangeEpk feedRange, CancellationToken cancellationToken = default)
         {
-            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            Routing.PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             IReadOnlyList<PartitionKeyRange> pkRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
                 collectionRid: await this.container.GetCachedRIDAsync(
                     forceRefresh: false,

diff --git a/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs
@@ -81,8 +81,8 @@ public override async Task<ResponseMessage> SendAsync(
 
                     DocumentServiceRequest serviceRequest = request.ToDocumentServiceRequest();
 
-                    PartitionKeyRangeCache routingMapProvider = await this.client.DocumentClient.GetPartitionKeyRangeCacheAsync();
-                    CollectionCache collectionCache = await this.client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton);
+                    PartitionKeyRangeCache routingMapProvider = await this.client.DocumentClient.GetPartitionKeyRangeCacheAsync(childTrace);
+                    CollectionCache collectionCache = await this.client.DocumentClient.GetCollectionCacheAsync(childTrace);
                     ContainerProperties collectionFromCache =
                         await collectionCache.ResolveCollectionAsync(serviceRequest, CancellationToken.None, childTrace);
 

diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs
@@ -70,7 +70,7 @@ public override async Task<ResponseMessage> SendAsync(
                 }
 
                 await this.ValidateAndSetConsistencyLevelAsync(request);
-                (bool isError, ResponseMessage errorResponse) = await this.EnsureValidClientAsync(request);
+                (bool isError, ResponseMessage errorResponse) = await this.EnsureValidClientAsync(request, childTrace);
                 if (isError)
                 {
                     return errorResponse;
@@ -196,7 +196,7 @@ public virtual async Task<ResponseMessage> SendAsync(
                         {
                             DocumentServiceRequest serviceRequest = request.ToDocumentServiceRequest();
 
-                            PartitionKeyRangeCache routingMapProvider = await this.client.DocumentClient.GetPartitionKeyRangeCacheAsync();
+                            PartitionKeyRangeCache routingMapProvider = await this.client.DocumentClient.GetPartitionKeyRangeCacheAsync(childTrace);
                             CollectionCache collectionCache = await this.client.DocumentClient.GetCollectionCacheAsync(childTrace);
                             ContainerProperties collectionFromCache =
                                 await collectionCache.ResolveCollectionAsync(serviceRequest, cancellationToken, childTrace);
@@ -323,11 +323,11 @@ internal static HttpMethod GetHttpMethod(
             }
         }
 
-        private async Task<(bool, ResponseMessage)> EnsureValidClientAsync(RequestMessage request)
+        private async Task<(bool, ResponseMessage)> EnsureValidClientAsync(RequestMessage request, ITrace trace)
         {
             try
             {
-                await this.client.DocumentClient.EnsureValidClientAsync();
+                await this.client.DocumentClient.EnsureValidClientAsync(trace);
                 return RequestInvokerHandler.clientIsValid;
             }
             catch (DocumentClientException dce)

diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs
@@ -219,7 +219,7 @@ internal async Task AssertPartitioningDetailsAsync(CosmosClient client, Cancella
 #if DEBUG
             try
             {
-                CollectionCache collectionCache = await client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton);
+                CollectionCache collectionCache = await client.DocumentClient.GetCollectionCacheAsync(trace);
                 ContainerProperties collectionFromCache =
                     await collectionCache.ResolveCollectionAsync(this.ToDocumentServiceRequest(), cancellationToken, trace);
                 if (collectionFromCache.PartitionKey?.Paths?.Count > 0)

diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs
@@ -82,7 +82,7 @@ async Task<CollectionCache> IDocumentQueryClient.GetCollectionCacheAsync()
 
         async Task<IRoutingMapProvider> IDocumentQueryClient.GetRoutingMapProviderAsync()
         {
-            return await this.innerClient.GetPartitionKeyRangeCacheAsync();
+            return await this.innerClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
         }
 
         public Task<QueryPartitionProvider> GetQueryPartitionProviderAsync()
@@ -118,7 +118,7 @@ public Task EnsureValidOverwriteAsync(ConsistencyLevel requestedConsistencyLevel
 
         public Task<PartitionKeyRangeCache> GetPartitionKeyRangeCacheAsync()
         {
-            return this.innerClient.GetPartitionKeyRangeCacheAsync();
+            return this.innerClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
         }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs
@@ -414,7 +414,7 @@ public override async Task<IReadOnlyList<PartitionKeyRange>> TryGetOverlappingRa
 
         private Task<PartitionKeyRangeCache> GetRoutingMapProviderAsync()
         {
-            return this.documentClient.GetPartitionKeyRangeCacheAsync();
+            return this.documentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
         }
 
         /// <summary>

diff --git a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs
@@ -260,7 +260,7 @@ public async Task<IReadOnlyList<FeedRange>> GetFeedRangesAsync(
             ITrace trace,
             CancellationToken cancellationToken = default)
         {
-            PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(trace);
 
             string containerRId;
             containerRId = await this.GetCachedRIDAsync(
@@ -381,7 +381,7 @@ internal async Task<IEnumerable<string>> GetPartitionKeyRangesAsync(
             ITrace trace,
             CancellationToken cancellationToken = default)
         {
-            IRoutingMapProvider routingMapProvider = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            IRoutingMapProvider routingMapProvider = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(trace);
             string containerRid = await this.GetCachedRIDAsync(
                 forceRefresh: false,
                 trace,
@@ -499,7 +499,7 @@ public override async Task<CollectionRoutingMap> GetRoutingMapAsync(Cancellation
                 trace: NoOpTrace.Singleton,
                 cancellationToken);
 
-            PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.Client.DocumentClient.GetPartitionKeyRangeCacheAsync();
+            PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.Client.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
             CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(
                 collectionRid,
                 previousValue: null,