add graceful handling of failures in QueryPhaseResultConsumer

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

Author: kaushalmahi12

server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java

@@ -217,6 +217,9 @@ private MergeResult partialReduce(
         int numReducePhases
     ) {
         checkCancellation();
+        if (pendingMerges.hasFailure()) {
+            return lastMerge;
+        }
         // ensure consistent ordering
         Arrays.sort(toConsume, Comparator.comparingInt(QuerySearchResult::getShardIndex));
 
@@ -450,6 +453,10 @@ private synchronized void onMergeFailure(Exception exc) {
             MergeTask task = runningTask.get();
             runningTask.compareAndSet(task, null);
             onPartialMergeFailure.accept(exc);
+            clearPendingMerges(task);
+        }
+
+        void clearPendingMerges(MergeTask task) {
             List<MergeTask> toCancels = new ArrayList<>();
             if (task != null) {
                 toCancels.add(task);
@@ -471,10 +478,11 @@ private void resetCircuitBreakerForCurrentRequest() {
 
         private void onAfterMerge(MergeTask task, MergeResult newResult, long estimatedSize) {
             synchronized (this) {
+                runningTask.compareAndSet(task, null);
                 if (hasFailure()) {
+                    task.cancel();
                     return;
                 }
-                runningTask.compareAndSet(task, null);
                 mergeResult = newResult;
                 if (hasAggs) {
                     // Update the circuit breaker to remove the size of the source aggregations
@@ -495,7 +503,11 @@ private void onAfterMerge(MergeTask task, MergeResult newResult, long estimatedS
         private void tryExecuteNext() {
             final MergeTask task;
             synchronized (this) {
-                if (queue.isEmpty() || hasFailure() || runningTask.get() != null) {
+                if (hasFailure()) {
+                    clearPendingMerges(null);
+                    return;
+                }
+                if (queue.isEmpty() || runningTask.get() != null) {
                     return;
                 }
                 task = queue.poll();
@@ -511,6 +523,7 @@ protected void doRun() {
                     try {
                         final QuerySearchResult[] toConsume = task.consumeBuffer();
                         if (toConsume == null) {
+                            task.cancel();
                             return;
                         }
                         long estimatedMergeSize = estimateRamBytesUsedForReduce(estimatedTotalSize);

server/src/test/java/org/opensearch/action/search/SearchPhaseControllerTests.java

@@ -1747,8 +1747,11 @@ public void testCancellationWithoutCircuitBreaker() throws Exception {
         int batchedReduceSize = randomIntBetween(2, expectedNumResults - 1);
         SearchRequest request = getAggregationSearchRequestWithBatchedReduceSize(batchedReduceSize);
         AssertingCircuitBreaker circuitBreaker = new AssertingCircuitBreaker(CircuitBreaker.REQUEST);
-        AtomicInteger checkCount = new AtomicInteger(0);
-        int cancelAfter = expectedNumResults / 2;
+        // To make it deterministic, we can count the number of times the partialReduce and reduce are called
+        // The exception is only thrown during the call to reduce which will happen once all shard level
+        // results have arrived
+        int partialReduceMethodCallCount = expectedNumResults / batchedReduceSize;
+        AtomicInteger checkCount = new AtomicInteger(expectedNumResults + partialReduceMethodCallCount);
 
         QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(
             fixedExecutor,
@@ -1758,7 +1761,7 @@ public void testCancellationWithoutCircuitBreaker() throws Exception {
             expectedNumResults,
             exc -> {},
             () -> {
-                return checkCount.incrementAndGet() > cancelAfter;
+                return checkCount.decrementAndGet() <= 0;
             }
         );
 
@@ -1775,9 +1778,8 @@ public void testCancellationDoesNotMaskCircuitBreakerException() throws Exceptio
 
         // making sure circuit breaker trips first
         circuitBreaker.shouldBreak.set(true);
-        AtomicInteger checkCount = new AtomicInteger(0);
-        int cancelAfter = expectedNumResults + 1;
-
+        int partialReduceMethodCallCount = expectedNumResults / batchedReduceSize;
+        AtomicInteger checkCount = new AtomicInteger(expectedNumResults + partialReduceMethodCallCount);
         QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(
             fixedExecutor,
             circuitBreaker,
@@ -1786,7 +1788,7 @@ public void testCancellationDoesNotMaskCircuitBreakerException() throws Exceptio
             expectedNumResults,
             exc -> {},
             () -> {
-                return checkCount.incrementAndGet() > cancelAfter;
+                return checkCount.decrementAndGet() <= 0;
             }
         );
 
@@ -1826,13 +1828,7 @@ private static void consumeShardLevelQueryPhaseResultsAsync(int expectedNumResul
                 result.setShardIndex(index);
                 result.size(1);
 
-                try {
-                    consumer.consumeResult(result, latch::countDown);
-                } catch (Exception e) {
-                    // Ensure latch counts down even on cancellation
-                    latch.countDown();
-                    // Don't rethrow - let the thread complete normally
-                }
+                consumer.consumeResult(result, latch::countDown);
             });
             threads[index].start();
         }