Fix: Hybrid Search Scheduling Stopwatch to Use Cumulative Timing Across All Phases (Azure#46591)

* Adding fixes for stopWatch when running hybrid search queries

* Updating changelog

* Resolving comments

* Resolving comments

Author: aayush3011

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java

@@ -130,8 +130,8 @@ public void hybridQueryTest() {
         resultDocs = container.queryItems(querySpec, new CosmosQueryRequestOptions(), Document.class).byPage()
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))
             .collectList().block();
-        assertThat(resultDocs).hasSize(13);
-        validateResults(Arrays.asList("51", "49", "61", "24", "54", "2", "22", "75", "77", "57", "76", "80", "85"), resultDocs);
+        assertThat(resultDocs).hasSize(15);
+        validateResults(Arrays.asList("51", "49", "24", "61", "54", "22", "2", "25", "75", "77", "57", "76", "66", "80", "85"), resultDocs);
 
         query = "SELECT c.id, c.title FROM c WHERE FullTextContains(c.title, @term1) " +
             "OR FullTextContains(c.text, @term1) OR FullTextContains(c.text, @term2) ORDER BY " +
@@ -143,8 +143,8 @@ public void hybridQueryTest() {
         resultDocs = container.queryItems(querySpec, new CosmosQueryRequestOptions(), Document.class).byPage()
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))
             .collectList().block();
-        assertThat(resultDocs).hasSize(8);
-        validateResults(Arrays.asList("2", "22", "75", "77", "57", "76", "80", "85"), resultDocs);
+        assertThat(resultDocs).hasSize(10);
+        validateResults(Arrays.asList("22", "2", "25", "75", "77", "57", "76", "66", "80", "85"), resultDocs);
 
         List<Float> vector = getQueryVector();
         query = "SELECT TOP 10 c.id, c.text, c.title FROM c " +
@@ -159,7 +159,7 @@ public void hybridQueryTest() {
             .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults()))
             .collectList().block();
         assertThat(resultDocs).hasSize(10);
-        validateResults(Arrays.asList("4", "24", "6", "9", "2", "3", "21", "5", "13", "49"), resultDocs);
+        validateResults(Arrays.asList("4", "24", "6", "9", "2", "3", "21", "5", "49", "13"), resultDocs);
     }
 
     @Test(groups = {"query", "split"}, timeOut = TIMEOUT)

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -12,6 +12,8 @@
 * Fixed Null Pointer Exception for query when container recreated with same name. - [PR 45930](https://github.com/Azure/azure-sdk-for-java/pull/45930)
 * Fixed Null Pointer Exception for readMany when container recreated with same name. - [PR 45930](https://github.com/Azure/azure-sdk-for-java/pull/45930)
 * Fixed parameterized query failures for Hybrid Search queries. - [PR 46446](https://github.com/Azure/azure-sdk-for-java/pull/46446)
+* Fixed a rare race in parallel Hybrid Search queries by making internal SchedulingStopwatch start/stop atomic and idempotent. - [PR 46485](https://github.com/Azure/azure-sdk-for-java/pull/46485)
+* Fixed the Hybrid Search SchedulingStopWatch to use Cumulative Timing Across all parallel queries. - [PR 46591](https://github.com/Azure/azure-sdk-for-java/pull/46591)
 * Fixed Strong Consistency violation when a single replica in a partition returns a 410 `Lease Not Found`. - [PR 46433](https://github.com/Azure/azure-sdk-for-java/pull/46433)
 
 #### Other Changes

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentProducer.java

@@ -7,6 +7,7 @@
 import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
+import com.azure.cosmos.implementation.query.metrics.SchedulingStopwatch;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
 import com.azure.cosmos.models.FeedResponse;
 import reactor.core.publisher.Mono;
@@ -16,6 +17,8 @@
 import java.util.function.Supplier;
 
 public class HybridSearchDocumentProducer extends DocumentProducer<Document> {
+    private final SchedulingStopwatch sharedSchedulingStopwatch;
+
     public HybridSearchDocumentProducer(
         IDocumentQueryClient client,
         String collectionResourceId,
@@ -30,9 +33,11 @@ public HybridSearchDocumentProducer(
         int initialPageSize,
         String initialContinuationToken,
         int top,
-        Supplier<String> operationContextTextProvider) {
+        Supplier<String> operationContextTextProvider,
+        SchedulingStopwatch sharedSchedulingStopwatch) {
         super(client, collectionResourceId, cosmosQueryRequestOptions, createRequestFunc, executeRequestFunc,
             collectionLink, createRetryPolicyFunc, resourceType, correlatedActivityId, initialPageSize,
             initialContinuationToken, top, feedRange, operationContextTextProvider);
+            this.sharedSchedulingStopwatch = sharedSchedulingStopwatch;
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java

@@ -15,6 +15,8 @@
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryResult;
+import com.azure.cosmos.implementation.query.metrics.HybridSearchCumulativeSchedulingStopWatch;
+import com.azure.cosmos.implementation.query.metrics.SchedulingStopwatch;
 import com.azure.cosmos.implementation.QueryMetrics;
 import com.azure.cosmos.implementation.RequestChargeTracker;
 import com.azure.cosmos.implementation.ResourceType;
@@ -69,6 +71,7 @@ public class HybridSearchDocumentQueryExecutionContext extends ParallelDocumentQ
     private final Collection<ClientSideRequestStatistics> clientSideRequestStatistics;
     private Flux<HybridSearchQueryResult<Document>> hybridObservable;
     private Mono<GlobalFullTextSearchQueryStatistics> aggregatedGlobalStatistics;
+    private final SchedulingStopwatch hybridSearchSchedulingStopwatch;
 
     protected HybridSearchDocumentQueryExecutionContext(
         DiagnosticsClientContext diagnosticsClientContext,
@@ -88,6 +91,10 @@ protected HybridSearchDocumentQueryExecutionContext(
         this.tracker = new RequestChargeTracker();
         this.queryMetricMap = new ConcurrentHashMap<>();
         this.clientSideRequestStatistics = ConcurrentHashMap.newKeySet();
+
+        // Initialize the shared stopwatch for hybrid search timing
+        this.hybridSearchSchedulingStopwatch = new HybridSearchCumulativeSchedulingStopWatch();
+        this.hybridSearchSchedulingStopwatch.ready();
     }
 
     public static Flux<IDocumentQueryExecutionComponent<Document>> createAsync(
@@ -130,6 +137,8 @@ private void initialize(
         int initialPageSize,
         DocumentCollection collection) {
 
+        // Start the hybrid search cumulative stopwatch when search begins
+        this.hybridSearchSchedulingStopwatch.start();
 
         if (hybridSearchQueryInfo.getRequiresGlobalStatistics()) {
             Map<FeedRangeEpkImpl, String> partitionKeyRangeToContinuationToken = new HashMap<>();
@@ -175,7 +184,11 @@ private Flux<HybridSearchQueryResult<Document>> hybridSearch(List<FeedRangeEpkIm
         Mono<List<List<Integer>>> ranks = computeRanks(componentScoresList);
 
         // Compute the RRF scores
-        return computeRRFScores(ranks, coalescedAndSortedResults, componentWeights);
+        return computeRRFScores(ranks, coalescedAndSortedResults, componentWeights)
+            .doFinally(signalType -> {
+                // Stop the hybrid search cumulative stopwatch when search completes
+                this.hybridSearchSchedulingStopwatch.stop();
+            });
     }
 
     @Override
@@ -206,7 +219,8 @@ protected HybridSearchDocumentProducer createDocumentProducer(
             initialPageSize,
             continuationToken,
             top,
-            this.getOperationContextTextProvider());
+            this.getOperationContextTextProvider(),
+            this.hybridSearchSchedulingStopwatch);
     }
 
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/metrics/HybridSearchCumulativeSchedulingStopWatch.java

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.implementation.query.metrics;
+
+import com.azure.cosmos.implementation.apachecommons.lang.time.StopWatch;
+
+public class HybridSearchCumulativeSchedulingStopWatch extends SchedulingStopwatch {
+    private long cumulativeRunTime;
+    private StopWatch runTimeStopwatch;
+    private long numPreemptions;
+    private boolean responded;
+
+    public HybridSearchCumulativeSchedulingStopWatch() {
+        super();
+        this.cumulativeRunTime = 0;
+        this.runTimeStopwatch = new StopWatch();
+    }
+
+    @Override
+    public SchedulingTimeSpan getElapsedTime() {
+        SchedulingTimeSpan parentTimeSpan = super.getElapsedTime();
+
+        long totalRunTime = this.cumulativeRunTime;
+        if (this.runTimeStopwatch.isStarted()) {
+            totalRunTime += this.runTimeStopwatch.getTime();
+        }
+
+        return new SchedulingTimeSpan(
+            parentTimeSpan.getTurnaroundTime(),
+            parentTimeSpan.getResponseTime(),
+            totalRunTime,
+            parentTimeSpan.getTurnaroundTime() - totalRunTime,
+            this.numPreemptions
+        );
+    }
+
+    @Override
+    public void start() {
+        synchronized (this.runTimeStopwatch) {
+            if (this.runTimeStopwatch.isStarted()) {
+                return;
+            }
+            if (!this.responded) {
+                this.responded = true;
+            }
+            // Don't reset - allow cumulative timing
+            this.runTimeStopwatch.start();
+        }
+    }
+
+    @Override
+    public void stop() {
+        synchronized (this.runTimeStopwatch) {
+            if (!this.runTimeStopwatch.isStarted()) {
+                return;
+            }
+            this.runTimeStopwatch.stop();
+            // Add elapsed time to cumulative total
+            this.cumulativeRunTime += this.runTimeStopwatch.getTime();
+            // Reset for next cycle
+            this.runTimeStopwatch.reset();
+            this.numPreemptions++;
+        }
+    }
+}