Diagnostics: Fixes GetQueryMetrics returning null inside custom RequestHandler

Microsoft.Azure.Cosmos/src/Handler/TransportHandler.cs

@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Handlers
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
     using Microsoft.Azure.Cosmos.Resource.CosmosExceptions;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Cosmos.Tracing.TraceData;
@@ -155,6 +156,29 @@ internal async Task<ResponseMessage> ProcessMessageAsync(
                     request,
                     serviceRequest.RequestContext.RequestChargeTracker);
 
+                // Attach query metrics to the per-request transport trace so they are
+                // reachable from a custom RequestHandler.SendAsync as the response unwinds
+                // back through the handler pipeline (issue #5117). Before this, the datum
+                // was added in CosmosQueryClientCore.GetCosmosElementResponse, which runs
+                // after the handler pipeline has fully unwound, so handlers always saw
+                // a null GetQueryMetrics result. Guarded on operation type to avoid
+                // mis-tagging non-query operations that may ever surface a query-metrics
+                // header. The thin-client / distributed-gateway path adds the same datum
+                // from CosmosDistributedQueryClient.CreatePage (that path bypasses
+                // TransportHandler).
+                if ((request.OperationType == OperationType.Query
+                        || request.OperationType == OperationType.SqlQuery)
+                    && !string.IsNullOrEmpty(responseMessage?.Headers?.QueryMetricsText))
+                {
+                    string queryMetricsText = responseMessage.Headers.QueryMetricsText;
+                    QueryMetricsTraceDatum queryMetricsDatum = new QueryMetricsTraceDatum(
+                        new Lazy<QueryMetrics>(() => new QueryMetrics(
+                            queryMetricsText,
+                            IndexUtilizationInfo.Empty,
+                            ClientSideMetrics.Empty)));
+                    processMessageAsyncTrace.AddDatum(TraceDatumKeys.QueryMetrics, queryMetricsDatum);
+                }
+
                 // Enrich diagnostics context in-case of auth failures 
                 if (responseMessage?.StatusCode == System.Net.HttpStatusCode.Unauthorized || responseMessage?.StatusCode == System.Net.HttpStatusCode.Forbidden)
                 {

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosDistributedQueryClient.cs

@@ -120,8 +120,11 @@ private DocumentServiceRequest CreateRequest(
             return request;
         }
 
-        private static TryCatch<QueryPage> CreatePage(DocumentServiceResponse response, Tracing.ITrace trace)
-        {
+        private static TryCatch<QueryPage> CreatePage(DocumentServiceResponse response, Tracing.ITrace trace)
+        {
+            // Attach QueryMetricsTraceDatum for the thin-client / distributed-gateway path
+            // (this path bypasses Handler/TransportHandler.cs which writes the same datum
+            // for the REST/RNTBD path). Keep both call sites in sync — see issue #5117.
             string queryMetricsText = response.Headers[HttpConstants.HttpHeaders.QueryMetrics];
             if (queryMetricsText != null)
             {

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -16,7 +16,6 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Json;
     using Microsoft.Azure.Cosmos.Query.Core;
-    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
     using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
@@ -322,15 +321,11 @@ private static TryCatch<QueryPage> GetCosmosElementResponse(
             {
                 using (cosmosResponseMessage)
                 {
-                    if (cosmosResponseMessage.Headers.QueryMetricsText != null)
-                    {
-                        QueryMetricsTraceDatum datum = new QueryMetricsTraceDatum(
-                            new Lazy<QueryMetrics>(() => new QueryMetrics(
-                                cosmosResponseMessage.Headers.QueryMetricsText, 
-                                IndexUtilizationInfo.Empty, 
-                                ClientSideMetrics.Empty)));
-                        trace.AddDatum(TraceDatumKeys.QueryMetrics, datum);
-                    }
+                    // Note: QueryMetricsTraceDatum is attached at the transport layer
+                    // (Handler/TransportHandler.cs) for the REST/RNTBD path so the datum
+                    // is reachable from custom RequestHandlers as the response unwinds.
+                    // The thin-client path attaches it in
+                    // Query/Core/QueryClient/CosmosDistributedQueryClient.CreatePage.
 
                     if (!cosmosResponseMessage.IsSuccessStatusCode)
                     {

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

@@ -5,7 +5,9 @@
 namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
 {
     using System;
+    using System.Collections.Concurrent;
     using System.Collections.Generic;
+    using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Documents;
@@ -95,6 +97,197 @@ public async Task TestBatchRequiredHeadersWithHandler()
             }
         }
 
+        [TestMethod]
+        public async Task QueryMetricsAvailableInsideRequestHandler_SinglePartition()
+        {
+            // Regression test for issue #5117:
+            // response.Diagnostics.GetQueryMetrics() must be non-null when invoked inside
+            // a custom RequestHandler.SendAsync for a query operation.
+            ConcurrentBag<ServerSideCumulativeMetrics> observedMetrics = new ConcurrentBag<ServerSideCumulativeMetrics>();
+
+            RequestHandlerHelper testHandler = new RequestHandlerHelper
+            {
+                CallBackOnResponse = (request, response) =>
+                {
+                    if (request.OperationType == OperationType.Query
+                        && response.IsSuccessStatusCode
+                        && !string.IsNullOrEmpty(response.Headers?.QueryMetricsText))
+                    {
+                        ServerSideCumulativeMetrics metrics = response.Diagnostics.GetQueryMetrics();
+                        Assert.IsNotNull(metrics, "GetQueryMetrics() must be non-null inside the handler (issue #5117).");
+                        observedMetrics.Add(metrics);
+                    }
+
+                    return response;
+                }
+            };
+
+            CosmosClient customClient = TestCommon.CreateCosmosClient(
+                builder => builder.AddCustomHandlers(testHandler));
+
+            try
+            {
+                Container container = customClient.GetContainer(this.database.Id, this.Container.Id);
+                QueryMetricsItem item = new QueryMetricsItem
+                {
+                    id = Guid.NewGuid().ToString(),
+                    pk = Guid.NewGuid().ToString(),
+                    description = "single-partition"
+                };
+                await container.CreateItemAsync(item, new Cosmos.PartitionKey(item.pk));
+
+                FeedIterator<QueryMetricsItem> iterator = container.GetItemQueryIterator<QueryMetricsItem>(
+                    $"SELECT * FROM c WHERE c.id = '{item.id}'",
+                    requestOptions: new QueryRequestOptions
+                    {
+                        PartitionKey = new Cosmos.PartitionKey(item.pk)
+                    });
+
+                int totalItems = 0;
+                while (iterator.HasMoreResults)
+                {
+                    FeedResponse<QueryMetricsItem> page = await iterator.ReadNextAsync();
+                    totalItems += page.Count;
+                }
+
+                Assert.AreEqual(1, totalItems);
+                Assert.IsTrue(
+                    observedMetrics.Count >= 1,
+                    "Handler must be invoked for at least one query response.");
+
+                foreach (ServerSideCumulativeMetrics metrics in observedMetrics)
+                {
+                    Assert.IsNotNull(metrics);
+                    Assert.AreEqual(
+                        1,
+                        metrics.PartitionedMetrics.Count,
+                        "Single-partition query should report exactly one PartitionedMetrics entry.");
+                }
+            }
+            finally
+            {
+                customClient.Dispose();
+            }
+        }
+
+        [TestMethod]
+        public async Task QueryMetricsAvailableInsideRequestHandler_CrossPartition()
+        {
+            // Companion to the single-partition test. Verifies that for cross-partition
+            // queries the handler still sees non-null metrics on every query response.
+            // Option C semantics: response.Diagnostics walks from the operation root, so
+            // each handler invocation sees its own partition's metrics plus any sibling
+            // partitions that have already completed at that point. The exact per-page
+            // distribution is timing-dependent; what we guarantee is (a) metrics are not
+            // null inside the handler and (b) the iterator-level aggregated view exposes
+            // metrics from multiple partitions.
+            ConcurrentBag<int> observedPartitionCounts = new ConcurrentBag<int>();
+
+            RequestHandlerHelper testHandler = new RequestHandlerHelper
+            {
+                CallBackOnResponse = (request, response) =>
+                {
+                    if (request.OperationType == OperationType.Query
+                        && response.IsSuccessStatusCode
+                        && !string.IsNullOrEmpty(response.Headers?.QueryMetricsText))
+                    {
+                        ServerSideCumulativeMetrics metrics = response.Diagnostics.GetQueryMetrics();
+                        Assert.IsNotNull(metrics, "GetQueryMetrics() must be non-null inside the handler (issue #5117).");
+                        observedPartitionCounts.Add(metrics.PartitionedMetrics.Count);
+                    }
+
+                    return response;
+                }
+            };
+
+            CosmosClient customClient = TestCommon.CreateCosmosClient(
+                builder => builder.AddCustomHandlers(testHandler));
+
+            Cosmos.Database multiPartitionDatabase = null;
+            try
+            {
+                multiPartitionDatabase = await customClient.CreateDatabaseAsync(
+                    "MultiPkDb_" + Guid.NewGuid().ToString());
+
+                Container multiPartitionContainer = await multiPartitionDatabase.CreateContainerAsync(
+                    new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: "/pk"),
+                    throughput: 15000);
+
+                IReadOnlyList<FeedRange> feedRanges = await multiPartitionContainer.GetFeedRangesAsync();
+                if (feedRanges.Count < 2)
+                {
+                    Assert.Inconclusive(
+                        $"Emulator did not split the 15000 RU container into >= 2 physical partitions (got {feedRanges.Count}); cannot exercise cross-partition fan-out.");
+                }
+
+                for (int i = 0; i < 30; i++)
+                {
+                    QueryMetricsItem item = new QueryMetricsItem
+                    {
+                        id = Guid.NewGuid().ToString(),
+                        pk = Guid.NewGuid().ToString(),
+                        description = "cross-partition-fanout"
+                    };
+                    await multiPartitionContainer.CreateItemAsync(item, new Cosmos.PartitionKey(item.pk));
+                }
+
+                FeedIterator<QueryMetricsItem> iterator = multiPartitionContainer.GetItemQueryIterator<QueryMetricsItem>(
+                    "SELECT * FROM c",
+                    requestOptions: new QueryRequestOptions
+                    {
+                        MaxConcurrency = -1
+                    });
+
+                int totalItems = 0;
+                int maxPartitionedCountOnIterator = 0;
+                bool iteratorSawMetrics = false;
+                while (iterator.HasMoreResults)
+                {
+                    FeedResponse<QueryMetricsItem> page = await iterator.ReadNextAsync();
+                    totalItems += page.Count;
+
+                    ServerSideCumulativeMetrics pageMetrics = page.Diagnostics.GetQueryMetrics();
+                    if (pageMetrics != null)
+                    {
+                        iteratorSawMetrics = true;
+                        if (pageMetrics.PartitionedMetrics.Count > maxPartitionedCountOnIterator)
+                        {
+                            maxPartitionedCountOnIterator = pageMetrics.PartitionedMetrics.Count;
+                        }
+                    }
+                }
+
+                Assert.AreEqual(30, totalItems);
+
+                Assert.IsTrue(
+                    observedPartitionCounts.Count >= 2,
+                    $"Cross-partition query should invoke the handler for at least 2 partition responses, got {observedPartitionCounts.Count}.");
+
+                foreach (int count in observedPartitionCounts)
+                {
+                    Assert.IsTrue(
+                        count >= 1,
+                        "Each handler invocation should see at least one partition's metrics (issue #5117).");
+                }
+
+                Assert.IsTrue(
+                    iteratorSawMetrics,
+                    "Expected at least one page's iterator-level diagnostics to expose query metrics.");
+                Assert.IsTrue(
+                    maxPartitionedCountOnIterator >= 2,
+                    $"Iterator-level diagnostics should aggregate metrics from multiple partitions on at least one page (max observed was {maxPartitionedCountOnIterator}).");
+            }
+            finally
+            {
+                if (multiPartitionDatabase != null)
+                {
+                    await multiPartitionDatabase.DeleteAsync();
+                }
+
+                customClient.Dispose();
+            }
+        }
+
         private async Task<IList<ToDoActivity>> CreateRandomItems(int pkCount, int perPKItemCount = 1, bool randomPartitionKey = true)
         {
             Assert.IsFalse(!randomPartitionKey && perPKItemCount > 1);
@@ -146,5 +339,12 @@ public class ToDoActivity
             public string description { get; set; }
             public string status { get; set; }
         }
+
+        private class QueryMetricsItem
+        {
+            public string id { get; set; }
+            public string pk { get; set; }
+            public string description { get; set; }
+        }
     }
 }

changelog.md

@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 #### Bugs Fixed
 
 - [5827](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/5827) ChangeFeedEstimator: Change feed estimator threw `ArgumentNullException` when an inmemory lease container was being used. Update validations so in-memory lease containers work with change feed estimator
+- [5879](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/5879) Diagnostics: Fixes `response.Diagnostics.GetQueryMetrics()` returning `null` when called inside a custom `RequestHandler.SendAsync` for query operations. Query metrics for the REST/RNTBD path are now attached at the transport layer, so they are visible to handlers as the response unwinds through the pipeline. `response.Diagnostics` continues to expose the full operation context (handlers see metrics from any sibling partition responses that have already completed). `FeedResponse.Diagnostics.GetQueryMetrics()` at the iterator level is unchanged.
 
 #### Other Changes