Query : Adds Request Charge to Query Metrics

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

@@ -111,7 +111,7 @@ private ReadOnlyMemory<byte> WriteTraceToJsonWriter(JsonSerializationFormat json
         private static ServerSideCumulativeMetrics PopulateServerSideCumulativeMetrics(ITrace trace)
         {
             ServerSideMetricsInternalAccumulator accumulator = new ServerSideMetricsInternalAccumulator();
-            ServerSideMetricsInternalAccumulator.WalkTraceTreeForQueryMetrics(trace, accumulator);
+            ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(trace, accumulator);
 
             IReadOnlyList<ServerSidePartitionedMetricsInternal> serverSideMetricsList = accumulator.GetPartitionedServerSideMetrics().Select(metrics => new ServerSidePartitionedMetricsInternal(metrics)).ToList();
 

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideCumulativeMetrics.cs

@@ -15,6 +15,11 @@ public abstract class ServerSideCumulativeMetrics
         /// </summary>
         public abstract ServerSideMetrics CumulativeMetrics { get; }
 
+        /// <summary>
+        /// Gets the total request charge for all partitions.
+        /// </summary>
+        public abstract double TotalRequestCharge { get; }
+
         /// <summary>
         /// Gets the list of ServerSideMetrics, one for for each partition.
         /// </summary>

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideCumulativeMetricsInternal.cs

@@ -20,10 +20,13 @@ internal ServerSideCumulativeMetricsInternal(IEnumerable<ServerSidePartitionedMe
         {
             this.PartitionedMetrics = serverSideMetricsList.ToList();
             this.CumulativeMetrics = ServerSideMetricsInternal.Create(serverSideMetricsList.Select(partitionedMetrics => partitionedMetrics.ServerSideMetricsInternal));
+            this.TotalRequestCharge = serverSideMetricsList.Sum(partitionedMetrics => partitionedMetrics.RequestCharge);
         }
 
         public override ServerSideMetrics CumulativeMetrics { get; }
 
         public override IReadOnlyList<ServerSidePartitionedMetrics> PartitionedMetrics { get; }
+
+        public override double TotalRequestCharge { get; }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsInternal.cs

@@ -124,6 +124,8 @@ sealed class ServerSideMetricsInternal : ServerSideMetrics
 
         public int? PartitionKeyRangeId { get; set; }
 
+        public double RequestCharge { get; set; }
+
         public static ServerSideMetricsInternal Create(IEnumerable<ServerSideMetricsInternal> serverSideMetricsEnumerable)
         {
             ServerSideMetricsInternalAccumulator accumulator = new ServerSideMetricsInternalAccumulator();

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsInternalAccumulator.cs

@@ -80,67 +80,5 @@ public List<ServerSideMetricsInternal> GetPartitionedServerSideMetrics()
         {
             return this.serverSideMetricsList;
         }
-
-        public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideMetricsInternalAccumulator accumulator)
-        {
-            if (currentTrace == null)
-            {
-                return;
-            }
-
-            foreach (object datum in currentTrace.Data.Values)
-            {
-                if (datum is QueryMetricsTraceDatum queryMetricsTraceDatum)
-                {
-                    queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.FeedRange = currentTrace.Name;
-                    queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.PartitionKeyRangeId = WalkTraceTreeForPartitionKeyRangeId(currentTrace);
-                    accumulator.Accumulate(queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics);
-                    return;
-                }
-            }
-
-            foreach (ITrace childTrace in currentTrace.Children)
-            {
-                WalkTraceTreeForQueryMetrics(childTrace, accumulator);
-            }
-
-            return;
-        }
-
-        private static int? WalkTraceTreeForPartitionKeyRangeId(ITrace currentTrace)
-        {
-            if (currentTrace == null)
-            {
-                return null;
-            }
-
-            foreach (Object datum in currentTrace.Data.Values)
-            {
-                if (datum is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)
-                {
-                    if (clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList.Count > 0)
-                    {
-                        return int.TryParse(clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList[0].StoreResult.PartitionKeyRangeId, out int pKRangeId)
-                            ? pKRangeId
-                            : null;
-                    }
-                    else
-                    {
-                        return null;
-                    }
-                }
-            }
-
-            foreach (ITrace childTrace in currentTrace.Children)
-            {
-                int? partitionKeyRangeId = WalkTraceTreeForPartitionKeyRangeId(childTrace);
-                if (partitionKeyRangeId != null)
-                {
-                    return partitionKeyRangeId;
-                }
-            }
-
-            return null;
-        }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsTraceExtractor.cs

@@ -0,0 +1,70 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.Metrics
+{
+    using System;
+    using Microsoft.Azure.Cosmos.Tracing;
+    using Microsoft.Azure.Cosmos.Tracing.TraceData;
+
+    internal static class ServerSideMetricsTraceExtractor
+    {
+        public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideMetricsInternalAccumulator accumulator)
+        {
+            if (currentTrace == null)
+            {
+                return;
+            }
+
+            foreach (object datum in currentTrace.Data.Values)
+            {
+                if (datum is QueryMetricsTraceDatum queryMetricsTraceDatum)
+                {
+                    ServerSideMetricsInternal serverSideMetrics = queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics;
+                    serverSideMetrics.FeedRange = currentTrace.Name;
+                    ServerSideMetricsTraceExtractor.WalkTraceTreeForPartitionInfo(currentTrace, serverSideMetrics);
+                    accumulator.Accumulate(serverSideMetrics);
+                }
+            }
+
+            foreach (ITrace childTrace in currentTrace.Children)
+            {
+                ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(childTrace, accumulator);
+            }
+        }
+
+        private static void WalkTraceTreeForPartitionInfo(ITrace currentTrace, ServerSideMetricsInternal serverSideMetrics)
+        {
+            if (currentTrace == null)
+            {
+                return;
+            }
+
+            foreach (Object datum in currentTrace.Data.Values)
+            {
+                if (datum is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)
+                {
+                    if (clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList.Count > 0)
+                    {
+                        if (int.TryParse(clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList[0].StoreResult.PartitionKeyRangeId, out int pKRangeId))
+                        {
+                            serverSideMetrics.PartitionKeyRangeId = pKRangeId;
+                        }
+
+                        serverSideMetrics.RequestCharge = clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList[0].StoreResult.RequestCharge;
+                    }
+                }
+                else if (datum is PointOperationStatisticsTraceDatum pointOperationStatisticsTraceDatum)
+                {
+                    serverSideMetrics.RequestCharge = pointOperationStatisticsTraceDatum.RequestCharge;
+                }
+            }
+
+            foreach (ITrace childTrace in currentTrace.Children)
+            {
+                ServerSideMetricsTraceExtractor.WalkTraceTreeForPartitionInfo(childTrace, serverSideMetrics);
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSidePartitionedMetrics.cs

@@ -26,5 +26,10 @@ public abstract class ServerSidePartitionedMetrics
         /// Only has a value in direct mode. When using gateway mode, this is null.
         /// </remarks>
         public abstract int? PartitionKeyRangeId { get; }
+
+        /// <summary>
+        /// Gets the request charge for the operation on this partition.
+        /// </summary>
+        public abstract double RequestCharge { get; }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSidePartitionedMetricsInternal.cs

@@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos
     internal class ServerSidePartitionedMetricsInternal : ServerSidePartitionedMetrics
     {
         internal ServerSidePartitionedMetricsInternal(ServerSideMetricsInternal serverSideMetricsInternal)
-            : this(serverSideMetricsInternal, serverSideMetricsInternal.FeedRange, serverSideMetricsInternal.PartitionKeyRangeId)
+            : this(serverSideMetricsInternal, serverSideMetricsInternal.FeedRange, serverSideMetricsInternal.PartitionKeyRangeId, serverSideMetricsInternal.RequestCharge)
         {
         }
 
@@ -22,11 +22,13 @@ internal ServerSidePartitionedMetricsInternal(ServerSideMetricsInternal serverSi
         /// <param name="serverSideMetricsInternal"></param>
         /// <param name="feedRange"></param>
         /// <param name="partitionKeyRangeId"></param>
-        internal ServerSidePartitionedMetricsInternal(ServerSideMetricsInternal serverSideMetricsInternal, string feedRange, int? partitionKeyRangeId)
+        /// <param name="requestCharge"></param>
+        internal ServerSidePartitionedMetricsInternal(ServerSideMetricsInternal serverSideMetricsInternal, string feedRange, int? partitionKeyRangeId, double requestCharge)
         {
             this.ServerSideMetricsInternal = serverSideMetricsInternal;
             this.FeedRange = feedRange;
             this.PartitionKeyRangeId = partitionKeyRangeId;
+            this.RequestCharge = requestCharge;
         }
 
         public ServerSideMetricsInternal ServerSideMetricsInternal { get; }
@@ -36,5 +38,7 @@ internal ServerSidePartitionedMetricsInternal(ServerSideMetricsInternal serverSi
         public override string FeedRange { get; }
 
         public override int? PartitionKeyRangeId { get; }
+
+        public override double RequestCharge { get; }
     }
 }

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

@@ -33,7 +33,6 @@ 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
@@ -1287,8 +1286,10 @@ public async Task QuerySinglePartitionItemStreamTest(int perPKItemCount, int max
 
                 ServerSideCumulativeMetrics metrics = response.Diagnostics.GetQueryMetrics();
                 Assert.IsTrue(metrics.PartitionedMetrics.Count > 0);
+                Assert.IsTrue(metrics.PartitionedMetrics[0].RequestCharge > 0);
                 Assert.IsTrue(metrics.CumulativeMetrics.TotalTime > TimeSpan.Zero);
                 Assert.IsTrue(metrics.CumulativeMetrics.QueryPreparationTime > TimeSpan.Zero);
+                Assert.IsTrue(metrics.TotalRequestCharge > 0);
 
                 if (metrics.CumulativeMetrics.RetrievedDocumentCount >= 1)
                 {
@@ -1376,11 +1377,14 @@ public async Task ItemMultiplePartitionQuery()
                     Assert.IsTrue(metrics.PartitionedMetrics.Count == 3);
                     Assert.IsTrue(metrics.CumulativeMetrics.TotalTime > TimeSpan.Zero);
                     Assert.IsTrue(metrics.CumulativeMetrics.QueryPreparationTime > TimeSpan.Zero);
+                    Assert.IsTrue(metrics.TotalRequestCharge > 0);
 
                     foreach (ServerSidePartitionedMetrics partitionedMetrics in metrics.PartitionedMetrics)
                     {
                         Assert.IsNotNull(partitionedMetrics);
+                        Assert.IsNotNull(partitionedMetrics.FeedRange);
                         Assert.IsNotNull(partitionedMetrics.PartitionKeyRangeId);
+                        Assert.IsTrue(partitionedMetrics.RequestCharge > 0);
                     }
 
                     if (metrics.CumulativeMetrics.RetrievedDocumentCount >= 1)
@@ -1455,11 +1459,14 @@ public async Task ItemSinglePartitionQueryGateway()
                     Assert.IsTrue(metrics.PartitionedMetrics.Count == 1);
                     Assert.IsTrue(metrics.CumulativeMetrics.TotalTime > TimeSpan.Zero);
                     Assert.IsTrue(metrics.CumulativeMetrics.QueryPreparationTime > TimeSpan.Zero);
+                    Assert.IsTrue(metrics.TotalRequestCharge > 0);
 
                     foreach (ServerSidePartitionedMetrics partitionedMetrics in metrics.PartitionedMetrics)
                     {
                         Assert.IsNotNull(partitionedMetrics);
+                        Assert.IsNotNull(partitionedMetrics.FeedRange);
                         Assert.IsNull(partitionedMetrics.PartitionKeyRangeId);
+                        Assert.IsTrue(partitionedMetrics.RequestCharge > 0);
                     }
 
                     if (metrics.CumulativeMetrics.RetrievedDocumentCount >= 1)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json

@@ -9235,6 +9235,16 @@
     "Microsoft.Azure.Cosmos.ServerSideCumulativeMetrics;System.Object;IsAbstract:True;IsSealed:False;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
       "Subclasses": {},
       "Members": {
+        "Double get_TotalRequestCharge()": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Double get_TotalRequestCharge();IsAbstract:True;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Double TotalRequestCharge": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "Double TotalRequestCharge;CanRead:True;CanWrite:False;Double get_TotalRequestCharge();IsAbstract:True;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Microsoft.Azure.Cosmos.ServerSideMetrics CumulativeMetrics": {
           "Type": "Property",
           "Attributes": [],
@@ -9387,6 +9397,16 @@
     "Microsoft.Azure.Cosmos.ServerSidePartitionedMetrics;System.Object;IsAbstract:True;IsSealed:False;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
       "Subclasses": {},
       "Members": {
+        "Double get_RequestCharge()": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Double get_RequestCharge();IsAbstract:True;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Double RequestCharge": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "Double RequestCharge;CanRead:True;CanWrite:False;Double get_RequestCharge();IsAbstract:True;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Microsoft.Azure.Cosmos.ServerSideMetrics get_ServerSideMetrics()": {
           "Type": "Method",
           "Attributes": [],