merged

Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs

@@ -224,7 +224,7 @@ private async Task<PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
             PartitionKeyRangeServerBatchRequest serverRequest,
             CancellationToken cancellationToken)
         {
-            CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create();
+            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();
             CosmosDiagnosticScope limiterScope = diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.Limiter");
             SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId);
             using (await limiter.UsingWaitAsync(cancellationToken))

Microsoft.Azure.Cosmos/src/Batch/TransactionalBatchOperationResult.cs

@@ -214,7 +214,7 @@ internal ResponseMessage ToResponseMessage()
                 requestMessage: null,
                 headers: headers,
                 cosmosException: null,
-                diagnostics: this.DiagnosticsContext ?? CosmosDiagnosticsContext.Create())
+                diagnostics: this.DiagnosticsContext ?? new CosmosDiagnosticsContextCore())
             {
                 Content = this.ResourceStream
             };

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -49,8 +49,6 @@ public ClientRetryPolicy(
             this.enableEndpointDiscovery = enableEndpointDiscovery;
             this.sessionTokenRetryCount = 0;
             this.canUseMultipleWriteLocations = false;
-
-            this.sharedStatistics = new CosmosClientSideRequestStatistics();
         }
 
         /// <summary> 
@@ -126,7 +124,19 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
             this.isReadRequest = request.IsReadOnlyRequest;
             this.canUseMultipleWriteLocations = this.globalEndpointManager.CanUseMultipleWriteLocations(request);
 
-            request.RequestContext.ClientRequestStatistics = this.sharedStatistics;
+            if (request.RequestContext.ClientRequestStatistics == null)
+            {
+                if (this.sharedStatistics == null)
+                {
+                    this.sharedStatistics = new CosmosClientSideRequestStatistics();
+                }
+
+                request.RequestContext.ClientRequestStatistics = this.sharedStatistics;
+            }
+            else
+            {
+                this.sharedStatistics = request.RequestContext.ClientRequestStatistics;
+            }
 
             // clear previous location-based routing directive
             request.RequestContext.ClearRouteToLocation();

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -493,21 +493,31 @@ public virtual Task<DatabaseResponse> CreateDatabaseIfNotExistsAsync(
                 // Doing a Read before Create will give us better latency for existing databases
                 DatabaseProperties databaseProperties = this.PrepareDatabaseProperties(id);
                 Database database = this.GetDatabase(id);
-                ResponseMessage response = await database.ReadStreamAsync(requestOptions: requestOptions, cancellationToken: cancellationToken);
-                if (response.StatusCode != HttpStatusCode.NotFound)
+                ResponseMessage readResponse = await database.ReadStreamAsync(
+                    requestOptions: requestOptions,
+                    cancellationToken: cancellationToken);
+
+                if (readResponse.StatusCode != HttpStatusCode.NotFound)
                 {
-                    return await this.ClientContext.ResponseFactory.CreateDatabaseResponseAsync(database, Task.FromResult(response));
+                    return await this.ClientContext.ResponseFactory.CreateDatabaseResponseAsync(database, Task.FromResult(readResponse));
                 }
 
-                response = await this.CreateDatabaseStreamAsync(databaseProperties, throughput, requestOptions, cancellationToken);
-                if (response.StatusCode != HttpStatusCode.Conflict)
+                ResponseMessage createResponse = await this.CreateDatabaseStreamAsync(databaseProperties, throughput, requestOptions, cancellationToken);
+
+                // Merge the diagnostics with the first read request.
+                createResponse.DiagnosticsContext.AddDiagnosticsInternal(readResponse.DiagnosticsContext);
+                if (createResponse.StatusCode != HttpStatusCode.Conflict)
                 {
-                    return await this.ClientContext.ResponseFactory.CreateDatabaseResponseAsync(this.GetDatabase(databaseProperties.Id), Task.FromResult(response));
+                    return await this.ClientContext.ResponseFactory.CreateDatabaseResponseAsync(this.GetDatabase(databaseProperties.Id), Task.FromResult(createResponse));
                 }
 
                 // This second Read is to handle the race condition when 2 or more threads have Read the database and only one succeeds with Create
                 // so for the remaining ones we should do a Read instead of throwing Conflict exception
-                return await database.ReadAsync(cancellationToken: cancellationToken);
+                ResponseMessage readResponseAfterConflict = await database.ReadStreamAsync(
+                    requestOptions: requestOptions,
+                    cancellationToken: cancellationToken);
+                readResponseAfterConflict.DiagnosticsContext.AddDiagnosticsInternal(readResponse.DiagnosticsContext);
+                return await this.ClientContext.ResponseFactory.CreateDatabaseResponseAsync(this.GetDatabase(databaseProperties.Id), Task.FromResult(readResponseAfterConflict));
             });
         }
 

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs

@@ -224,6 +224,17 @@ public static CosmosElement Parse(string json)
 
             return cosmosElement;
         }
+
+        public static TCosmosElement Parse<TCosmosElement>(string json)
+            where TCosmosElement : CosmosElement
+        {
+            if (!CosmosElement.TryParse(json, out TCosmosElement cosmosElement))
+            {
+                throw new ArgumentException($"Failed to parse json: {json}.");
+            }
+
+            return cosmosElement;
+        }
     }
 #if INTERNAL
 #pragma warning restore SA1600 // Elements should be documented

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.LazyCosmosObject.cs

@@ -4,6 +4,7 @@
 namespace Microsoft.Azure.Cosmos.CosmosElements
 {
     using System;
+    using System.Collections.Concurrent;
     using System.Collections.Generic;
     using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
@@ -21,7 +22,7 @@ private class LazyCosmosObject : CosmosObject
         {
             private readonly IJsonNavigator jsonNavigator;
             private readonly IJsonNavigatorNode jsonNavigatorNode;
-            private readonly Dictionary<string, CosmosElement> cachedElements;
+            private readonly ConcurrentDictionary<string, CosmosElement> cachedElements;
             private readonly Lazy<int> lazyCount;
 
             public LazyCosmosObject(IJsonNavigator jsonNavigator, IJsonNavigatorNode jsonNavigatorNode)
@@ -44,7 +45,7 @@ public LazyCosmosObject(IJsonNavigator jsonNavigator, IJsonNavigatorNode jsonNav
 
                 this.jsonNavigator = jsonNavigator;
                 this.jsonNavigatorNode = jsonNavigatorNode;
-                this.cachedElements = new Dictionary<string, CosmosElement>();
+                this.cachedElements = new ConcurrentDictionary<string, CosmosElement>();
                 this.lazyCount = new Lazy<int>(() => this.jsonNavigator.GetObjectPropertyCount(this.jsonNavigatorNode));
             }
 
@@ -92,31 +93,27 @@ public override IEnumerator<KeyValuePair<string, CosmosElement>> GetEnumerator()
 
             public override bool TryGetValue(string key, out CosmosElement value)
             {
-                value = default;
-                bool gotValue;
                 if (this.cachedElements.TryGetValue(
                     key,
                     out CosmosElement cosmosElemet))
                 {
                     value = cosmosElemet;
-                    gotValue = true;
+                    return true;
                 }
-                else if (this.jsonNavigator.TryGetObjectProperty(
+
+                if (this.jsonNavigator.TryGetObjectProperty(
                     this.jsonNavigatorNode,
                     key,
                     out ObjectProperty objectProperty))
                 {
                     value = CosmosElement.Dispatch(this.jsonNavigator, objectProperty.ValueNode);
-                    gotValue = true;
                     this.cachedElements[key] = value;
-                }
-                else
-                {
-                    value = null;
-                    gotValue = false;
+
+                    return true;
                 }
 
-                return gotValue;
+                value = default;
+                return false;
             }
 
             public override void WriteTo(IJsonWriter jsonWriter)

Microsoft.Azure.Cosmos/src/Diagnostics/AddressResolutionStatistics.cs

@@ -0,0 +1,34 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+namespace Microsoft.Azure.Cosmos.Diagnostics
+{
+    using System;
+
+    internal sealed class AddressResolutionStatistics : CosmosDiagnosticsInternal
+    {
+        public AddressResolutionStatistics(
+            DateTime startTime,
+            DateTime endTime,
+            string targetEndpoint)
+        {
+            this.StartTime = startTime;
+            this.EndTime = endTime;
+            this.TargetEndpoint = targetEndpoint ?? throw new ArgumentNullException(nameof(startTime));
+        }
+
+        public DateTime StartTime { get; }
+        public DateTime? EndTime { get; set; }
+        public string TargetEndpoint { get; }
+
+        public override void Accept(CosmosDiagnosticsInternalVisitor visitor)
+        {
+            visitor.Visit(this);
+        }
+
+        public override TResult Accept<TResult>(CosmosDiagnosticsInternalVisitor<TResult> visitor)
+        {
+            return visitor.Visit(this);
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Diagnostics/BackendMetricsExtractor.cs

@@ -13,6 +13,7 @@ namespace Microsoft.Azure.Cosmos.Diagnostics
     internal sealed class BackendMetricsExtractor : CosmosDiagnosticsInternalVisitor<(ParseFailureReason, BackendMetrics)>
     {
         public static readonly BackendMetricsExtractor Singleton = new BackendMetricsExtractor();
+        private static readonly (ParseFailureReason, BackendMetrics) MetricsNotFound = (ParseFailureReason.MetricsNotFound, default);
 
         private BackendMetricsExtractor()
         {
@@ -21,7 +22,7 @@ private BackendMetricsExtractor()
 
         public override (ParseFailureReason, BackendMetrics) Visit(PointOperationStatistics pointOperationStatistics)
         {
-            return (ParseFailureReason.MetricsNotFound, default);
+            return BackendMetricsExtractor.MetricsNotFound;
         }
 
         public override (ParseFailureReason, BackendMetrics) Visit(CosmosDiagnosticsContext cosmosDiagnosticsContext)
@@ -61,7 +62,7 @@ public override (ParseFailureReason, BackendMetrics) Visit(CosmosDiagnosticsCont
 
         public override (ParseFailureReason, BackendMetrics) Visit(CosmosDiagnosticScope cosmosDiagnosticScope)
         {
-            return (ParseFailureReason.MetricsNotFound, default);
+            return BackendMetricsExtractor.MetricsNotFound;
         }
 
         public override (ParseFailureReason, BackendMetrics) Visit(QueryPageDiagnostics queryPageDiagnostics)
@@ -74,6 +75,21 @@ public override (ParseFailureReason, BackendMetrics) Visit(QueryPageDiagnostics
             return (ParseFailureReason.None, backendMetrics);
         }
 
+        public override (ParseFailureReason, BackendMetrics) Visit(AddressResolutionStatistics addressResolutionStatistics)
+        {
+            return BackendMetricsExtractor.MetricsNotFound;
+        }
+
+        public override (ParseFailureReason, BackendMetrics) Visit(StoreResponseStatistics storeResponseStatistics)
+        {
+            return BackendMetricsExtractor.MetricsNotFound;
+        }
+
+        public override (ParseFailureReason, BackendMetrics) Visit(CosmosClientSideRequestStatistics clientSideRequestStatistics)
+        {
+            return BackendMetricsExtractor.MetricsNotFound;
+        }
+
         public enum ParseFailureReason
         {
             None,