From a3eba570db66a1ba1697561d02811a510bfbe71d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Cea=20Fontenla?= Date: Thu, 7 Nov 2024 11:21:53 +0100 Subject: [PATCH] Aggs: Add real memory CB call when building internal aggregators in buckets (#116329) Related with https://github.com/elastic/elasticsearch/issues/88128 This PR pretends to reduce the potential OOMs received when building internal aggregations. --- .../bucket/BucketsAggregator.java | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java index 7c7233b0eaa1d..e6c26c4278807 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java @@ -81,12 +81,7 @@ public final void collectBucket(LeafBucketCollector subCollector, int doc, long grow(bucketOrd + 1); int docCount = docCountProvider.getDocCount(doc); if (docCounts.increment(bucketOrd, docCount) == docCount) { - // We call the circuit breaker the time to time in order to give it a chance to check available - // memory in the parent breaker and break the execution if we are running out. To achieve that we - // are passing 0 as the estimated bytes every 1024 calls - if ((++callCount & 0x3FF) == 0) { - breaker.addEstimateBytesAndMaybeBreak(0, "allocated_buckets"); - } + updateCircuitBreaker("allocated_buckets"); } subCollector.collect(doc, bucketOrd); } @@ -179,6 +174,7 @@ protected final IntFunction buildSubAggsForBuckets(long[] prepareSubAggs(bucketOrdsToCollect); InternalAggregation[][] aggregations = new InternalAggregation[subAggregators.length][]; for (int i = 0; i < subAggregators.length; i++) { + updateCircuitBreaker("building_sub_aggregation"); aggregations[i] = subAggregators[i].buildAggregations(bucketOrdsToCollect); } return subAggsForBucketFunction(aggregations); @@ -415,4 +411,15 @@ protected void preGetSubLeafCollectors(LeafReaderContext ctx) throws IOException // Set LeafReaderContext to the doc_count provider docCountProvider.setLeafReaderContext(ctx); } + + /** + * This method calls the circuit breaker from time to time in order to give it a chance to check available + * memory in the parent breaker (Which should be a real memory breaker) and break the execution if we are running out. + * To achieve that, we are passing 0 as the estimated bytes every 1024 calls + */ + private void updateCircuitBreaker(String label) { + if ((++callCount & 0x3FF) == 0) { + breaker.addEstimateBytesAndMaybeBreak(0, label); + } + } }