Revise the concurrent pattern

Signed-off-by: bowenlan-amzn <[email protected]>

Author: bowenlan-amzn

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

@@ -302,7 +302,7 @@ class PendingMerges implements Releasable {
         private volatile long maxAggsCurrentBufferSize = 0;
 
         private final ArrayDeque<MergeTask> queue = new ArrayDeque<>();
-        private final AtomicReference<MergeTask> runningTask = new AtomicReference<>();
+        private final AtomicReference<MergeTask> runningTask = new AtomicReference<>(); // ensure only one task is running
         private final AtomicReference<Exception> failure = new AtomicReference<>();
 
         private final SearchPhaseController.TopDocsStats topDocsStats;
@@ -355,7 +355,7 @@ private synchronized long addEstimateAndMaybeBreak(long estimatedSize) {
             return circuitBreakerBytes;
         }
 
-        private void resetCircuitBreaker() {
+        private synchronized void resetCircuitBreaker() {
             if (circuitBreakerBytes > 0) {
                 circuitBreaker.addWithoutBreaking(-circuitBreakerBytes);
                 circuitBreakerBytes = 0;
@@ -388,12 +388,12 @@ private long estimateRamBytesUsedForReduce(long size) {
         void consume(QuerySearchResult result, Runnable callback) {
             checkCancellation();
 
-            if (processResult(result, callback)) {
+            if (consumeResult(result, callback)) {
                 callback.run();
             }
         }
 
-        private synchronized boolean processResult(QuerySearchResult result, Runnable callback) {
+        private synchronized boolean consumeResult(QuerySearchResult result, Runnable callback) {
             if (hasFailure()) {
                 result.consumeAll(); // release memory
                 return true;
@@ -455,8 +455,7 @@ protected void doRun() {
                     try {
                         final QuerySearchResult[] toConsume = task.consumeBuffer();
                         if (toConsume == null) {
-                            onAfterMerge(task, thisMergeResult, estimatedTotalSize);
-                            executor.execute(() -> tryExecuteNext());
+                            onAfterMerge(task, null, 0);
                             return;
                         }
                         long estimateRamBytesUsedForReduce = estimateRamBytesUsedForReduce(estimatedTotalSize);
@@ -469,7 +468,6 @@ protected void doRun() {
                         return;
                     }
                     onAfterMerge(task, newMerge, estimatedTotalSize);
-                    executor.execute(() -> tryExecuteNext());
                 }
 
                 @Override
@@ -480,29 +478,32 @@ public void onFailure(Exception exc) {
         }
 
         private void onAfterMerge(MergeTask task, MergeResult newResult, long estimatedSize) {
-            synchronized (this) {
-                runningTask.compareAndSet(task, null);
-                if (hasFailure()) {
-                    return;
-                }
-                mergeResult = newResult;
-                if (hasAggs) {
-                    // Update the circuit breaker to remove the size of the source aggregations
-                    // and replace the estimation with the serialized size of the newly reduced result.
-                    long newSize = mergeResult.estimatedSize - estimatedSize;
-                    addWithoutBreaking(newSize);
-                    logger.trace(
-                        "aggs partial reduction [{}->{}] max [{}]",
-                        estimatedSize,
-                        mergeResult.estimatedSize,
-                        maxAggsCurrentBufferSize
-                    );
+            if (newResult != null) {
+                synchronized (this) {
+                    if (hasFailure()) {
+                        return;
+                    }
+                    runningTask.compareAndSet(task, null);
+                    mergeResult = newResult;
+                    if (hasAggs) {
+                        // Update the circuit breaker to remove the size of the source aggregations
+                        // and replace the estimation with the serialized size of the newly reduced result.
+                        long newSize = mergeResult.estimatedSize - estimatedSize;
+                        addWithoutBreaking(newSize);
+                        logger.trace(
+                            "aggs partial reduction [{}->{}] max [{}]",
+                            estimatedSize,
+                            mergeResult.estimatedSize,
+                            maxAggsCurrentBufferSize
+                        );
+                    }
                 }
-                task.consumeListener();
             }
+            task.consumeListener();
+            executor.execute(this::tryExecuteNext);
         }
 
-        // Idempotent failure handling logic
+        // Idempotent and thread-safe failure handling
         private synchronized void onMergeFailure(Exception exc) {
             if (hasFailure()) {
                 assert circuitBreakerBytes == 0;
@@ -515,7 +516,7 @@ private synchronized void onMergeFailure(Exception exc) {
             cancelTaskOnFailure.accept(exc);
         }
 
-        private void clearMergeTaskQueue() {
+        private synchronized void clearMergeTaskQueue() {
             MergeTask task = runningTask.get();
             runningTask.compareAndSet(task, null);
             List<MergeTask> toCancels = new ArrayList<>();