Skip to content

Commit a3eba57

Browse files
authored
Aggs: Add real memory CB call when building internal aggregators in buckets (#116329)
Related with #88128 This PR pretends to reduce the potential OOMs received when building internal aggregations.
1 parent 4dc5afc commit a3eba57

File tree

1 file changed

+13
-6
lines changed

1 file changed

+13
-6
lines changed

server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,7 @@ public final void collectBucket(LeafBucketCollector subCollector, int doc, long
8181
grow(bucketOrd + 1);
8282
int docCount = docCountProvider.getDocCount(doc);
8383
if (docCounts.increment(bucketOrd, docCount) == docCount) {
84-
// We call the circuit breaker the time to time in order to give it a chance to check available
85-
// memory in the parent breaker and break the execution if we are running out. To achieve that we
86-
// are passing 0 as the estimated bytes every 1024 calls
87-
if ((++callCount & 0x3FF) == 0) {
88-
breaker.addEstimateBytesAndMaybeBreak(0, "allocated_buckets");
89-
}
84+
updateCircuitBreaker("allocated_buckets");
9085
}
9186
subCollector.collect(doc, bucketOrd);
9287
}
@@ -179,6 +174,7 @@ protected final IntFunction<InternalAggregations> buildSubAggsForBuckets(long[]
179174
prepareSubAggs(bucketOrdsToCollect);
180175
InternalAggregation[][] aggregations = new InternalAggregation[subAggregators.length][];
181176
for (int i = 0; i < subAggregators.length; i++) {
177+
updateCircuitBreaker("building_sub_aggregation");
182178
aggregations[i] = subAggregators[i].buildAggregations(bucketOrdsToCollect);
183179
}
184180
return subAggsForBucketFunction(aggregations);
@@ -415,4 +411,15 @@ protected void preGetSubLeafCollectors(LeafReaderContext ctx) throws IOException
415411
// Set LeafReaderContext to the doc_count provider
416412
docCountProvider.setLeafReaderContext(ctx);
417413
}
414+
415+
/**
416+
* This method calls the circuit breaker from time to time in order to give it a chance to check available
417+
* memory in the parent breaker (Which should be a real memory breaker) and break the execution if we are running out.
418+
* To achieve that, we are passing 0 as the estimated bytes every 1024 calls
419+
*/
420+
private void updateCircuitBreaker(String label) {
421+
if ((++callCount & 0x3FF) == 0) {
422+
breaker.addEstimateBytesAndMaybeBreak(0, label);
423+
}
424+
}
418425
}

0 commit comments

Comments
 (0)