support gateway

Azure · Jan 15, 2024 · 4aa0a48 · 4aa0a48
1 parent 9f20691
commit 4aa0a48
Showing 1 changed file with 31 additions and 1 deletion.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsInternalAccumulator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsInternalAccumulator.cs
@@ -94,7 +94,10 @@ public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideM
                 {
                     queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.FeedRange = currentTrace.Name;
                     queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.PartitionKeyRangeId = WalkTraceTreeForPartitionKeyRangeId(currentTrace);
-                    queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.RequestCharge = WalkTraceTreeForRequestCharge(currentTrace);
+                    queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.RequestCharge = queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics.PartitionKeyRangeId != null
+                        ? WalkTraceTreeForRequestCharge(currentTrace)
+                        : WalkTraceTreeForRequestChargeGateway(currentTrace);
+
                     accumulator.Accumulate(queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics);
                     return;
                 }
@@ -177,5 +180,32 @@ public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideM
 
             return null;
         }
+
+        private static double? WalkTraceTreeForRequestChargeGateway(ITrace currentTrace)
+        {
+            if (currentTrace == null)
+            {
+                return null;
+            }
+
+            foreach (Object datum in currentTrace.Data.Values)
+            {
+                if (datum is PointOperationStatisticsTraceDatum pointOperationStatisticsTraceDatum)
+                {
+                    return pointOperationStatisticsTraceDatum.RequestCharge;
+                }
+            }
+
+            foreach (ITrace childTrace in currentTrace.Children)
+            {
+                double? requestCharge = WalkTraceTreeForRequestChargeGateway(childTrace);
+                if (requestCharge != null)
+                {
+                    return requestCharge;
+                }
+            }
+
+            return null;
+        }
     }
 }