Pr comments

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

@@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Diagnostics
 
     internal sealed class CosmosTraceDiagnostics : CosmosDiagnostics
     {
+        private ServerSideMetrics cachedServerSideMetrics { get; set; }
         public CosmosTraceDiagnostics(ITrace trace)
         {
             if (trace == null)
@@ -30,6 +31,7 @@ public CosmosTraceDiagnostics(ITrace trace)
             }
 
             this.Value = rootTrace;
+            this.cachedServerSideMetrics = null;
         }
 
         public ITrace Value { get; }
@@ -51,9 +53,15 @@ public override TimeSpan GetClientElapsedTime()
 
         public override ServerSideMetrics GetQueryMetrics()
         {
+            if (this.cachedServerSideMetrics != null)
+            {
+                return this.cachedServerSideMetrics;
+            }
+
             ServerSideMetricsAccumulator accumulator = new ServerSideMetricsAccumulator();
             ServerSideMetricsAccumulator.WalkTraceTreeForQueryMetrics(this.Value, accumulator);
-            return new ServerSideMetrics(accumulator.GetServerSideMetrics());
+            this.cachedServerSideMetrics = new ServerSideMetrics(accumulator.GetServerSideMetrics());
+            return this.cachedServerSideMetrics;
         }
 
         internal bool IsGoneExceptionHit()

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs

@@ -33,6 +33,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     using System.Reflection;
     using System.Text.RegularExpressions;
     using Microsoft.Azure.Cosmos.Diagnostics;
+    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
 
     [TestClass]
     public class CosmosItemTests : BaseCosmosClientHelper
@@ -1261,6 +1262,7 @@ public async Task QuerySinglePartitionItemStreamTest(int perPKItemCount, int max
             do
             {
                 iterationCount++;
+                ServerSideMetricsAccumulator headerMetricsAccumulator = new ServerSideMetricsAccumulator();
                 FeedIterator feedIterator = this.Container.GetItemQueryStreamIterator(
                     sql,
                     continuationToken: lastContinuationToken,
@@ -1278,6 +1280,13 @@ public async Task QuerySinglePartitionItemStreamTest(int perPKItemCount, int max
                 System.Diagnostics.Trace.TraceInformation($"ContinuationToken: {lastContinuationToken}");
                 Newtonsoft.Json.JsonSerializer serializer = new Newtonsoft.Json.JsonSerializer();
 
+                // verify ServerSideMetrics matches metrics retrieved from header
+                ServerSideMetrics serverSideMetricsFromDiagnostics = response.Diagnostics.GetQueryMetrics();
+                bool tryParseResult = ServerSideMetricsInternal.TryParseFromDelimitedString(response.Headers.QueryMetricsText, out ServerSideMetricsInternal serverSideMetricsFromHeaders);
+                Assert.IsTrue(tryParseResult);
+                headerMetricsAccumulator.Accumulate(serverSideMetricsFromHeaders);
+                Assert.IsTrue(headerMetricsAccumulator.GetServerSideMetrics().FormatTrace() == serverSideMetricsFromDiagnostics.FormatTrace());
+
                 using (StreamReader sr = new StreamReader(response.Content))
                 using (JsonTextReader jtr = new JsonTextReader(sr))
                 {