[Backport 2.19] Harden the circuit breaker and failure handle logic in query result consumer

[jainankitk]

Description

Backports #19396 for preventing estimated circuit breaker limits from becoming negative. This issue was reported on one of the cluster for Amazon Opensearch Service

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit afc07bb.

'Diff too large, requires skip by maintainers after manual review'

---

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

---

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 78f7e42)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Circuit Breaker Accounting In the new consumeResult method, the circuit breaker is charged via addEstimateAndMaybeBreak before the result is added to the buffer, but aggsCurrentBufferSize is updated after. If a partial reduce task is triggered in the same call, aggsCurrentBufferSize is reset to 0 and the result is added to the buffer without its agg size being reflected. This could lead to under-accounting of memory in the circuit breaker for the buffered result. private synchronized boolean consumeResult(QuerySearchResult result, Runnable callback) { if (hasFailure()) { result.consumeAll(); // release memory return true; } if (result.isNull()) { SearchShardTarget target = result.getSearchShardTarget(); emptyResults.add(new SearchShard(target.getClusterAlias(), target.getShardId())); return true; } // Check circuit breaker before consuming if (hasAggs) { long aggsSize = ramBytesUsedQueryResult(result); try { addEstimateAndMaybeBreak(aggsSize); aggsCurrentBufferSize += aggsSize; } catch (CircuitBreakingException e) { onFailure(e); return true; } } // Process non-empty results int size = buffer.size() + (hasPartialReduce ? 1 : 0); if (size >= batchReduceSize) { hasPartialReduce = true; // the callback must wait for the new reduce task to complete to maintain proper result processing order QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); aggsCurrentBufferSize = 0; buffer.clear(); emptyResults.clear(); queue.add(task); tryExecuteNext(); buffer.add(result); return false; // callback will be run by reduce task } buffer.add(result); return true; } Estimation Change The estimateRamBytesUsedForReduce method now returns 0.5 * size instead of the previous 1.5 * size - size (which also equals 0.5 * size). While mathematically equivalent, the comment above still says "roughly 1.5 times the size", which is now misleading and inconsistent with the actual multiplier used. private long estimateRamBytesUsedForReduce(long size) { return Math.round(0.5d * size); } onAfterReduce Task Handling In onAfterReduce, when newResult is null (e.g., buffer was already consumed/cancelled), runningTask is never cleared via compareAndSet. This means hasPendingReduceTask() could return true indefinitely, blocking future reduce tasks from executing. private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimatedSize) { if (newResult != null) { synchronized (this) { if (hasFailure()) { return; } runningTask.compareAndSet(task, null); reduceResult = 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 = reduceResult.estimatedSize - estimatedSize; addWithoutBreaking(newSize); logger.trace( "aggs partial reduction [{}->{}] max [{}]", estimatedSize, reduceResult.estimatedSize, maxAggsCurrentBufferSize ); } } } task.consumeListener(); executor.execute(this::tryExecuteNext); } checkCancellation Side Effect checkCancellation now calls pendingReduces.onFailure(...) which invokes cancelTaskOnFailure.accept(exc) — meaning a cancellation check during partialReduce will trigger the external cancel callback. Previously, cancellation was handled more carefully to avoid masking circuit breaker exceptions. Verify that calling cancelTaskOnFailure on task cancellation during partial reduce does not cause double-cancellation or unexpected behavior. private void checkCancellation() { if (isTaskCancelled.getAsBoolean()) { pendingReduces.onFailure(new TaskCancelledException("request has been terminated")); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 78f7e42

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Fix running task not cleared on null result When newResult is null (i.e., the buffer was already consumed/cancelled), runningTask is never cleared, which means hasPendingReduceTask() will keep returning true and tryExecuteNext() will never schedule subsequent tasks. The runningTask.compareAndSet(task, null) call should be executed regardless of whether newResult is null. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [496-520] private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimatedSize) { - if (newResult != null) { - synchronized (this) { - if (hasFailure()) { - return; + synchronized (this) { + runningTask.compareAndSet(task, null); + if (newResult != null && !hasFailure()) { + reduceResult = newResult; + if (hasAggs) { + long newSize = reduceResult.estimatedSize - estimatedSize; + addWithoutBreaking(newSize); + logger.trace( + "aggs partial reduction [{}->{}] max [{}]", + estimatedSize, + reduceResult.estimatedSize, + maxAggsCurrentBufferSize + ); } - runningTask.compareAndSet(task, null); - reduceResult = newResult; - ... } } task.consumeListener(); executor.execute(this::tryExecuteNext); } Suggestion importance[1-10]: 8 __ Why: When newResult is null (task was cancelled/consumed), runningTask is never cleared via compareAndSet, which would cause hasPendingReduceTask() to return true indefinitely and block subsequent tasks from executing. This is a real correctness bug in the new code.	Medium
	Fix buffer size accounting after partial reduce trigger When a partial reduce is triggered, the current result is added to the buffer after the old buffer is drained into the ReduceTask. However, the aggs size for this new result was already accounted for in the circuit breaker via addEstimateAndMaybeBreak(aggsSize) above, but aggsCurrentBufferSize is reset to 0 and the new result's size is not re-added to aggsCurrentBufferSize. This means the next ReduceTask's aggsBufferSize will be 0 instead of the size of the current result, causing incorrect circuit breaker accounting. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [412-450] -private synchronized boolean consumeResult(QuerySearchResult result, Runnable callback) { - ... - // Process non-empty results - int size = buffer.size() + (hasPartialReduce ? 1 : 0); - if (size >= batchReduceSize) { - hasPartialReduce = true; - // the callback must wait for the new reduce task to complete to maintain proper result processing order - QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); - ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); - aggsCurrentBufferSize = 0; - buffer.clear(); - emptyResults.clear(); - queue.add(task); - tryExecuteNext(); - buffer.add(result); - return false; // callback will be run by reduce task +if (size >= batchReduceSize) { + hasPartialReduce = true; + QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); + ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); + aggsCurrentBufferSize = 0; + buffer.clear(); + emptyResults.clear(); + queue.add(task); + tryExecuteNext(); + buffer.add(result); + if (hasAggs) { + long aggsSize = ramBytesUsedQueryResult(result); + aggsCurrentBufferSize += aggsSize; } - buffer.add(result); - return true; + return false; } Suggestion importance[1-10]: 7 __ Why: After a partial reduce is triggered, aggsCurrentBufferSize is reset to 0 but the current result's agg size (already accounted in the circuit breaker) is not added back to aggsCurrentBufferSize. This causes the next ReduceTask's aggsBufferSize to be 0, leading to incorrect circuit breaker accounting for subsequent partial reduces.	Medium
General	Avoid holding lock during external callback invocation The cancelTaskOnFailure.accept(exc) is called while holding the synchronized lock on this. If cancelTaskOnFailure tries to acquire any lock that is also held by a thread waiting on this object's monitor, it could cause a deadlock. The external callback should be invoked outside the synchronized block. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [523-533] -private synchronized void onFailure(Exception exc) { - if (hasFailure()) { - assert circuitBreakerBytes == 0; - return; +private void onFailure(Exception exc) { + synchronized (this) { + if (hasFailure()) { + assert circuitBreakerBytes == 0; + return; + } + assert circuitBreakerBytes >= 0; + resetCircuitBreaker(); + failure.compareAndSet(null, exc); + clearReduceTaskQueue(); } - assert circuitBreakerBytes >= 0; - resetCircuitBreaker(); - failure.compareAndSet(null, exc); - clearReduceTaskQueue(); cancelTaskOnFailure.accept(exc); } Suggestion importance[1-10]: 7 __ Why: Calling cancelTaskOnFailure.accept(exc) while holding the synchronized lock on this risks deadlock if the callback tries to acquire any lock held by a thread waiting on this monitor. Moving the callback outside the synchronized block is a valid concurrency improvement.	Medium

---

Previous suggestions

Suggestions up to commit 21fcf4d

Category	Suggestion	Impact
Possible issue	Clear running task even when result is null When newResult is null (i.e., the buffer was already consumed/cancelled), runningTask is never cleared, which means hasPendingReduceTask() will keep returning true and tryExecuteNext() will never schedule further tasks. The runningTask.compareAndSet(task, null) call should be executed regardless of whether newResult is null. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [496-520] private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimatedSize) { - if (newResult != null) { - synchronized (this) { - if (hasFailure()) { - return; + synchronized (this) { + runningTask.compareAndSet(task, null); + if (newResult != null && !hasFailure()) { + reduceResult = newResult; + if (hasAggs) { + long newSize = reduceResult.estimatedSize - estimatedSize; + addWithoutBreaking(newSize); + logger.trace( + "aggs partial reduction [{}->{}] max [{}]", + estimatedSize, + reduceResult.estimatedSize, + maxAggsCurrentBufferSize + ); } - runningTask.compareAndSet(task, null); - reduceResult = newResult; - ... } } task.consumeListener(); executor.execute(this::tryExecuteNext); } Suggestion importance[1-10]: 8 __ Why: When newResult is null (buffer already consumed/cancelled), runningTask is never cleared, causing hasPendingReduceTask() to remain true and blocking further task scheduling. This is a real bug that could cause the system to stall. The fix correctly moves runningTask.compareAndSet(task, null) outside the null check.	Medium
	Fix buffer size tracking after partial reduce When a new ReduceTask is created and the current result is added to the buffer after clearing, the aggs size for this new result has already been accounted via addEstimateAndMaybeBreak(aggsSize) earlier in the method, but aggsCurrentBufferSize was reset to 0 before adding the result to the buffer. The aggsCurrentBufferSize should be updated to reflect the size of the newly buffered result after the reset. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [412-450] -private synchronized boolean consumeResult(QuerySearchResult result, Runnable callback) { - if (hasFailure()) { - result.consumeAll(); // release memory - return true; +if (size >= batchReduceSize) { + hasPartialReduce = true; + QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); + ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); + aggsCurrentBufferSize = 0; + buffer.clear(); + emptyResults.clear(); + queue.add(task); + tryExecuteNext(); + buffer.add(result); + if (hasAggs) { + long aggsSize = ramBytesUsedQueryResult(result); + aggsCurrentBufferSize += aggsSize; } - ... - // Process non-empty results - int size = buffer.size() + (hasPartialReduce ? 1 : 0); - if (size >= batchReduceSize) { - hasPartialReduce = true; - // the callback must wait for the new reduce task to complete to maintain proper result processing order - QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); - ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); - aggsCurrentBufferSize = 0; - buffer.clear(); - emptyResults.clear(); - queue.add(task); - tryExecuteNext(); - buffer.add(result); - return false; // callback will be run by reduce task - } - buffer.add(result); - return true; + return false; } Suggestion importance[1-10]: 7 __ Why: After a partial reduce is triggered, aggsCurrentBufferSize is reset to 0 but the newly added result's agg size (already accounted in the circuit breaker via addEstimateAndMaybeBreak) is not reflected in aggsCurrentBufferSize. This could lead to incorrect size tracking for subsequent reduce operations.	Medium
General	Avoid redundant synchronization in failure cleanup clearReduceTaskQueue() is also synchronized, so calling it from within onFailure() which is already synchronized will cause a deadlock if the lock is not reentrant. In Java, synchronized on instance methods uses the same intrinsic lock and is reentrant, so this is safe — but clearReduceTaskQueue should not be a separate synchronized method if it's always called from within a synchronized context, to avoid confusion and potential issues if the locking strategy changes. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [523-533] private synchronized void onFailure(Exception exc) { if (hasFailure()) { assert circuitBreakerBytes == 0; return; } assert circuitBreakerBytes >= 0; resetCircuitBreaker(); failure.compareAndSet(null, exc); - clearReduceTaskQueue(); + clearReduceTaskQueueLocked(); // private non-synchronized helper cancelTaskOnFailure.accept(exc); } +private void clearReduceTaskQueueLocked() { + ReduceTask task = runningTask.get(); + runningTask.compareAndSet(task, null); + List<ReduceTask> toCancels = new ArrayList<>(); + if (task != null) { + toCancels.add(task); + } + toCancels.addAll(queue); + queue.clear(); + reduceResult = null; + for (ReduceTask toCancel : toCancels) { + toCancel.cancel(); + } +} + Suggestion importance[1-10]: 3 __ Why: Java's intrinsic locks are reentrant, so calling a synchronized method from within another synchronized method on the same instance is safe and won't deadlock. This is a style/clarity suggestion rather than a correctness fix, and the concern raised is not actually a real issue in Java.	Low

Suggestions up to commit 9849d38

Category	Suggestion	Impact
Possible issue	Clear running task even when result is null When newResult is null (i.e., the buffer was already consumed/cancelled), runningTask is never cleared. This means hasPendingReduceTask() will keep returning true indefinitely, blocking the final reduce() call. The runningTask should be cleared regardless of whether newResult is null. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [492-516] private void onAfterReduce(ReduceTask task, ReduceResult newResult, long estimatedSize) { - if (newResult != null) { - synchronized (this) { - if (hasFailure()) { - return; + synchronized (this) { + if (!hasFailure()) { + runningTask.compareAndSet(task, null); + if (newResult != null) { + reduceResult = newResult; + if (hasAggs) { + long newSize = reduceResult.estimatedSize - estimatedSize; + addWithoutBreaking(newSize); + logger.trace( + "aggs partial reduction [{}->{}] max [{}]", + estimatedSize, + reduceResult.estimatedSize, + maxAggsCurrentBufferSize + ); + } } - runningTask.compareAndSet(task, null); - reduceResult = newResult; - ... } } task.consumeListener(); executor.execute(this::tryExecuteNext); } Suggestion importance[1-10]: 8 __ Why: When newResult is null (buffer already consumed/cancelled), runningTask is never cleared, causing hasPendingReduceTask() to return true indefinitely and blocking the final reduce() call. This is a real bug that could cause hangs.	Medium
	Fix buffer size tracking after partial reduce When a new ReduceTask is created and the current result is added to the buffer after clearing, the aggs size for this new result has already been accounted via addEstimateAndMaybeBreak(aggsSize) earlier, but aggsCurrentBufferSize was reset to 0. The new result's aggsSize should be re-added to aggsCurrentBufferSize after the reset so the next reduce task has an accurate buffer size estimate. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [431-443] -private synchronized boolean consumeResult(QuerySearchResult result, Runnable callback) { - if (hasFailure()) { - result.consumeAll(); // release memory - return true; +if (size >= batchReduceSize) { + hasPartialReduce = true; + QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); + ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); + aggsCurrentBufferSize = 0; + buffer.clear(); + emptyResults.clear(); + queue.add(task); + tryExecuteNext(); + buffer.add(result); + if (hasAggs) { + aggsCurrentBufferSize += ramBytesUsedQueryResult(result); } - ... - // Process non-empty results - int size = buffer.size() + (hasPartialReduce ? 1 : 0); - if (size >= batchReduceSize) { - hasPartialReduce = true; - // the callback must wait for the new reduce task to complete to maintain proper result processing order - QuerySearchResult[] clone = buffer.toArray(QuerySearchResult[]::new); - ReduceTask task = new ReduceTask(clone, aggsCurrentBufferSize, new ArrayList<>(emptyResults), callback); - aggsCurrentBufferSize = 0; - buffer.clear(); - emptyResults.clear(); - queue.add(task); - tryExecuteNext(); - buffer.add(result); - return false; // callback will be run by reduce task - } - buffer.add(result); - return true; + return false; } Suggestion importance[1-10]: 7 __ Why: After aggsCurrentBufferSize is reset to 0 and the current result is added to the buffer, the new result's agg size is not re-added to aggsCurrentBufferSize, causing inaccurate size estimates for subsequent reduce tasks. This is a valid correctness issue.	Medium
General	Avoid nested synchronized calls on same monitor clearReduceTaskQueue() is also synchronized, so calling it from within onFailure() (which is already synchronized) will cause a deadlock if the lock is not reentrant. In Java, synchronized on instance methods uses the same monitor, so re-entrant calls on the same thread are safe — however, clearReduceTaskQueue should not be a separate synchronized method if it's always called from within another synchronized block to avoid confusion and potential issues if the locking strategy changes. server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java [519-529] private synchronized void onFailure(Exception exc) { if (hasFailure()) { assert circuitBreakerBytes == 0; return; } assert circuitBreakerBytes >= 0; resetCircuitBreaker(); failure.compareAndSet(null, exc); - clearReduceTaskQueue(); + clearReduceTaskQueueLocked(); // private non-synchronized helper cancelTaskOnFailure.accept(exc); } +private void clearReduceTaskQueueLocked() { + // must be called while holding 'this' monitor + ReduceTask task = runningTask.get(); + runningTask.compareAndSet(task, null); + List<ReduceTask> toCancels = new ArrayList<>(); + if (task != null) { + toCancels.add(task); + } + toCancels.addAll(queue); + queue.clear(); + reduceResult = null; + for (ReduceTask toCancel : toCancels) { + toCancel.cancel(); + } +} + Suggestion importance[1-10]: 3 __ Why: Java's synchronized is reentrant on the same thread, so there is no actual deadlock risk here. The suggestion is about code style and future-proofing, which is a minor concern with limited practical impact.	Low

[jainankitk]

@bowenlan-amzn - Seems the record is not supported for 2.x. I don't see any harm in getting rid of record for backport 2.19 version. Unless you feel that change might be complex, and do more harm than good

> Task :server:compileJava
/home/runner/work/OpenSearch/OpenSearch/server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java:589: error: records are not supported in -source 11
    private record ReduceResult(List<SearchShard> processedShards, TopDocs reducedTopDocs, InternalAggregations reducedAggs,
            ^
  (use -source 16 or higher to enable records)
/home/runner/work/OpenSearch/OpenSearch/server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java:589: warning: 'record' may become a restricted type name in a future release and may be unusable for type declarations or as the element type of an array
    private record ReduceResult(List<SearchShard> processedShards, TopDocs reducedTopDocs, InternalAggregations reducedAggs,

[github-actions]

❌ Gradle check result for 9849d38: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 21fcf4d

[github-actions]

Persistent review updated to latest commit 78f7e42

[github-actions]

✅ Gradle check result for 78f7e42: SUCCESS

Code change is no go for 2.19.5

[codecov]

Codecov Report

❌ Patch coverage is 83.59375% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.02%. Comparing base (f915333) to head (78f7e42).
⚠️ Report is 9 commits behind head on 2.19.

Files with missing lines	Patch %	Lines
...search/action/search/QueryPhaseResultConsumer.java	83.59%	10 Missing and 11 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               2.19   #20769      +/-   ##
============================================
+ Coverage     71.97%   72.02%   +0.05%     
- Complexity    65995    66018      +23     
============================================
  Files          5342     5342              
  Lines        307363   307361       -2     
  Branches      44857    44857              
============================================
+ Hits         221211   221375     +164     
+ Misses        67661    67484     -177     
- Partials      18491    18502      +11     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.