Check the real memory circuit breaker when building internal aggregations (#117019) (#117247)

checks periodically the real memory circuit breaker when allocating objects.

modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/adjacency/AdjacencyMatrixAggregator.java

@@ -188,17 +188,16 @@ public class AdjacencyMatrixAggregator extends BucketsAggregator {
             }
         }
         try (LongArray bucketOrdsToBuild = bigArrays().newLongArray(totalBucketsToBuild)) {
-            int builtBucketIndex = 0;
+            int[] builtBucketIndex = new int[] { 0 };
             for (int ord = 0; ord < maxOrd; ord++) {
                 if (bucketDocCount(ord) > 0) {
-                    bucketOrdsToBuild.set(builtBucketIndex++, ord);
+                    bucketOrdsToBuild.set(builtBucketIndex[0]++, ord);
                 }
             }
-            assert builtBucketIndex == totalBucketsToBuild;
-            builtBucketIndex = 0;
+            assert builtBucketIndex[0] == totalBucketsToBuild;
+            builtBucketIndex[0] = 0;
             var bucketSubAggs = buildSubAggsForBuckets(bucketOrdsToBuild);
-            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-            for (int owningBucketOrdIdx = 0; owningBucketOrdIdx < results.length; owningBucketOrdIdx++) {
+            InternalAggregation[] aggregations = buildAggregations(Math.toIntExact(owningBucketOrds.size()), owningBucketOrdIdx -> {
                 List<InternalAdjacencyMatrix.InternalBucket> buckets = new ArrayList<>(filters.length);
                 for (int i = 0; i < keys.length; i++) {
                     long bucketOrd = bucketOrd(owningBucketOrds.get(owningBucketOrdIdx), i);
@@ -207,10 +206,11 @@ public class AdjacencyMatrixAggregator extends BucketsAggregator {
                     // a date-histogram where we will look for transactions over time and can expect many
                     // empty buckets.
                     if (docCount > 0) {
+                        checkRealMemoryCBForInternalBucket();
                         InternalAdjacencyMatrix.InternalBucket bucket = new InternalAdjacencyMatrix.InternalBucket(
                             keys[i],
                             docCount,
-                            bucketSubAggs.apply(builtBucketIndex++)
+                            bucketSubAggs.apply(builtBucketIndex[0]++)
                         );
                         buckets.add(bucket);
                     }
@@ -226,17 +226,17 @@ public class AdjacencyMatrixAggregator extends BucketsAggregator {
                             InternalAdjacencyMatrix.InternalBucket bucket = new InternalAdjacencyMatrix.InternalBucket(
                                 intersectKey,
                                 docCount,
-                                bucketSubAggs.apply(builtBucketIndex++)
+                                bucketSubAggs.apply(builtBucketIndex[0]++)
                             );
                             buckets.add(bucket);
                         }
                         pos++;
                     }
                 }
-                results[owningBucketOrdIdx] = new InternalAdjacencyMatrix(name, buckets, metadata());
-            }
-            assert builtBucketIndex == totalBucketsToBuild;
-            return results;
+                return new InternalAdjacencyMatrix(name, buckets, metadata());
+            });
+            assert builtBucketIndex[0] == totalBucketsToBuild;
+            return aggregations;
         }
     }
 

modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/TimeSeriesAggregator.java

@@ -79,6 +79,7 @@ public class TimeSeriesAggregator extends BucketsAggregator {
                 while (ordsEnum.next()) {
                     long docCount = bucketDocCount(ordsEnum.ord());
                     ordsEnum.readValue(spare);
+                    checkRealMemoryCBForInternalBucket();
                     InternalTimeSeries.InternalBucket bucket = new InternalTimeSeries.InternalBucket(
                         BytesRef.deepCopyOf(spare), // Closing bucketOrds will corrupt the bytes ref, so need to make a deep copy here.
                         docCount,
@@ -101,11 +102,7 @@ public class TimeSeriesAggregator extends BucketsAggregator {
             }
             buildSubAggsForAllBuckets(allBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
 
-            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(allBucketsPerOrd.size())];
-            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
-                result[ordIdx] = buildResult(allBucketsPerOrd.get(ordIdx));
-            }
-            return result;
+            return buildAggregations(Math.toIntExact(allBucketsPerOrd.size()), ordIdx -> buildResult(allBucketsPerOrd.get(ordIdx)));
         }
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java

@@ -13,6 +13,8 @@ import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.CheckedIntFunction;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.Maps;
@@ -48,6 +50,8 @@ public abstract class AggregatorBase extends Aggregator {
 
     private Map<String, Aggregator> subAggregatorbyName;
     private long requestBytesUsed;
+    private final CircuitBreaker breaker;
+    private int callCount;
 
     /**
      * Constructs a new Aggregator.
@@ -72,6 +76,7 @@ public abstract class AggregatorBase extends Aggregator {
         this.metadata = metadata;
         this.parent = parent;
         this.context = context;
+        this.breaker = context.breaker();
         assert factories != null : "sub-factories provided to BucketAggregator must not be null, use AggragatorFactories.EMPTY instead";
         this.subAggregators = factories.createSubAggregators(this, subAggregatorCardinality);
         context.addReleasable(this);
@@ -327,6 +332,30 @@ public abstract class AggregatorBase extends Aggregator {
         return InternalAggregations.from(aggs);
     }
 
+    /**
+     * Builds the aggregations array with the provided size and populates it using the provided function.
+     */
+    protected final InternalAggregation[] buildAggregations(int size, CheckedIntFunction<InternalAggregation, IOException> aggFunction)
+        throws IOException {
+        final InternalAggregation[] results = new InternalAggregation[size];
+        for (int i = 0; i < results.length; i++) {
+            checkRealMemoryCB("internal_aggregation");
+            results[i] = aggFunction.apply(i);
+        }
+        return results;
+    }
+
+    /**
+     * 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
+     */
+    protected final void checkRealMemoryCB(String label) {
+        if ((++callCount & 0x3FF) == 0) {
+            breaker.addEstimateBytesAndMaybeBreak(0, label);
+        }
+    }
+
     @Override
     public String toString() {
         return name;

server/src/main/java/org/elasticsearch/search/aggregations/NonCollectingAggregator.java

@@ -41,10 +41,6 @@ public abstract class NonCollectingAggregator extends AggregatorBase {
 
     @Override
     public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
-        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-            results[ordIdx] = buildEmptyAggregation();
-        }
-        return results;
+        return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> buildEmptyAggregation());
     }
 }

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

@@ -9,7 +9,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.util.IntArray;
 import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.ObjectArray;
@@ -42,10 +41,9 @@ import java.util.function.LongUnaryOperator;
 import java.util.function.ToLongFunction;
 
 public abstract class BucketsAggregator extends AggregatorBase {
-    private final CircuitBreaker breaker;
+
     private LongArray docCounts;
     protected final DocCountProvider docCountProvider;
-    private int callCount;
 
     @SuppressWarnings("this-escape")
     public BucketsAggregator(
@@ -57,7 +55,6 @@ public abstract class BucketsAggregator extends AggregatorBase {
         Map<String, Object> metadata
     ) throws IOException {
         super(name, factories, aggCtx, parent, bucketCardinality, metadata);
-        breaker = aggCtx.breaker();
         docCounts = bigArrays().newLongArray(1, true);
         docCountProvider = new DocCountProvider();
     }
@@ -83,7 +80,7 @@ public abstract class BucketsAggregator extends AggregatorBase {
         grow(bucketOrd + 1);
         int docCount = docCountProvider.getDocCount(doc);
         if (docCounts.increment(bucketOrd, docCount) == docCount) {
-            updateCircuitBreaker("allocated_buckets");
+            checkRealMemoryCB("allocated_buckets");
         }
         subCollector.collect(doc, bucketOrd);
     }
@@ -176,7 +173,7 @@ public abstract class BucketsAggregator extends AggregatorBase {
         prepareSubAggs(bucketOrdsToCollect);
         InternalAggregation[][] aggregations = new InternalAggregation[subAggregators.length][];
         for (int i = 0; i < subAggregators.length; i++) {
-            updateCircuitBreaker("building_sub_aggregation");
+            checkRealMemoryCB("building_sub_aggregation");
             aggregations[i] = subAggregators[i].buildAggregations(bucketOrdsToCollect);
         }
         return subAggsForBucketFunction(aggregations);
@@ -247,31 +244,30 @@ public abstract class BucketsAggregator extends AggregatorBase {
         Function<List<B>, InternalAggregation> resultBuilder
     ) throws IOException {
         try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(owningBucketOrds.size() * bucketsPerOwningBucketOrd)) {
-            int bucketOrdIdx = 0;
+            final int[] bucketOrdIdx = new int[] { 0 };
             for (long i = 0; i < owningBucketOrds.size(); i++) {
                 long ord = owningBucketOrds.get(i) * bucketsPerOwningBucketOrd;
                 for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
-                    bucketOrdsToCollect.set(bucketOrdIdx++, ord++);
+                    bucketOrdsToCollect.set(bucketOrdIdx[0]++, ord++);
                 }
             }
-            bucketOrdIdx = 0;
+            bucketOrdIdx[0] = 0;
             var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
 
-            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-            for (int owningOrdIdx = 0; owningOrdIdx < results.length; owningOrdIdx++) {
+            return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                 List<B> buckets = new ArrayList<>(bucketsPerOwningBucketOrd);
                 for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
+                    checkRealMemoryCBForInternalBucket();
                     buckets.add(
                         bucketBuilder.build(
                             offsetInOwningOrd,
-                            bucketDocCount(bucketOrdsToCollect.get(bucketOrdIdx)),
-                            subAggregationResults.apply(bucketOrdIdx++)
+                            bucketDocCount(bucketOrdsToCollect.get(bucketOrdIdx[0])),
+                            subAggregationResults.apply(bucketOrdIdx[0]++)
                         )
                     );
                 }
-                results[owningOrdIdx] = resultBuilder.apply(buckets);
-            }
-            return results;
+                return resultBuilder.apply(buckets);
+            });
         }
     }
 
@@ -295,11 +291,10 @@ public abstract class BucketsAggregator extends AggregatorBase {
          * here but we don't because single bucket aggs never have.
          */
         var subAggregationResults = buildSubAggsForBuckets(owningBucketOrds);
-        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-            results[ordIdx] = resultBuilder.build(owningBucketOrds.get(ordIdx), subAggregationResults.apply(ordIdx));
-        }
-        return results;
+        return buildAggregations(
+            Math.toIntExact(owningBucketOrds.size()),
+            ordIdx -> resultBuilder.build(owningBucketOrds.get(ordIdx), subAggregationResults.apply(ordIdx))
+        );
     }
 
     @FunctionalInterface
@@ -335,37 +330,36 @@ public abstract class BucketsAggregator extends AggregatorBase {
                 );
             }
             try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(totalOrdsToCollect)) {
-                int b = 0;
+                final int[] b = new int[] { 0 };
                 for (long i = 0; i < owningBucketOrds.size(); i++) {
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(i));
                     while (ordsEnum.next()) {
-                        bucketOrdsToCollect.set(b++, ordsEnum.ord());
+                        bucketOrdsToCollect.set(b[0]++, ordsEnum.ord());
                     }
                 }
                 var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
 
-                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-                b = 0;
-                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                b[0] = 0;
+                return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                     final long owningBucketOrd = owningBucketOrds.get(ordIdx);
                     List<B> buckets = new ArrayList<>(bucketsInOrd.get(ordIdx));
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
                     while (ordsEnum.next()) {
-                        if (bucketOrdsToCollect.get(b) != ordsEnum.ord()) {
+                        if (bucketOrdsToCollect.get(b[0]) != ordsEnum.ord()) {
                             // If we hit this, something has gone horribly wrong and we need to investigate
                             throw AggregationErrors.iterationOrderChangedWithoutMutating(
                                 bucketOrds.toString(),
                                 ordsEnum.ord(),
-                                bucketOrdsToCollect.get(b)
+                                bucketOrdsToCollect.get(b[0])
                             );
                         }
+                        checkRealMemoryCBForInternalBucket();
                         buckets.add(
-                            bucketBuilder.build(ordsEnum.value(), bucketDocCount(ordsEnum.ord()), subAggregationResults.apply(b++))
+                            bucketBuilder.build(ordsEnum.value(), bucketDocCount(ordsEnum.ord()), subAggregationResults.apply(b[0]++))
                         );
                     }
-                    results[ordIdx] = resultBuilder.build(owningBucketOrd, buckets);
-                }
-                return results;
+                    return resultBuilder.build(owningBucketOrd, buckets);
+                });
             }
         }
     }
@@ -425,14 +419,9 @@ public abstract class BucketsAggregator extends AggregatorBase {
         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);
-        }
+    /** This method should be called whenever a new bucket object is created. It will check the real memory
+     * circuit breaker in a sampling fashion. See {@link #checkRealMemoryCB(String)} */
+    protected final void checkRealMemoryCBForInternalBucket() {
+        checkRealMemoryCB("internal_bucket");
     }
 }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/countedterms/CountedTermsAggregator.java

@@ -140,6 +140,7 @@ class CountedTermsAggregator extends TermsAggregator {
                         long docCount = bucketDocCount(ordsEnum.ord());
                         otherDocCounts.increment(ordIdx, docCount);
                         if (spare == null) {
+                            checkRealMemoryCBForInternalBucket();
                             spare = emptyBucketBuilder.get();
                         }
                         ordsEnum.readValue(spare.getTermBytes());
@@ -158,8 +159,8 @@ class CountedTermsAggregator extends TermsAggregator {
             }
 
             buildSubAggsForAllBuckets(topBucketsPerOrd, InternalTerms.Bucket::getBucketOrd, InternalTerms.Bucket::setAggregations);
-            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
-            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+
+            return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                 final BucketOrder reduceOrder;
                 if (isKeyOrder(order) == false) {
                     reduceOrder = InternalOrder.key(true);
@@ -167,7 +168,7 @@ class CountedTermsAggregator extends TermsAggregator {
                 } else {
                     reduceOrder = order;
                 }
-                result[ordIdx] = new StringTerms(
+                return new StringTerms(
                     name,
                     reduceOrder,
                     order,
@@ -181,8 +182,7 @@ class CountedTermsAggregator extends TermsAggregator {
                     Arrays.asList(topBucketsPerOrd.get(ordIdx)),
                     null
                 );
-            }
-            return result;
+            });
         }
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregator.java

@@ -144,6 +144,7 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid<?>> extends Bu
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
                     while (ordsEnum.next()) {
                         if (spare == null) {
+                            checkRealMemoryCBForInternalBucket();
                             spare = newEmptyBucket();
                         }
 
@@ -162,11 +163,10 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid<?>> extends Bu
                 }
             }
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
-            for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-                results[ordIdx] = buildAggregation(name, requiredSize, Arrays.asList(topBucketsPerOrd.get(ordIdx)), metadata());
-            }
-            return results;
+            return buildAggregations(
+                Math.toIntExact(owningBucketOrds.size()),
+                ordIdx -> buildAggregation(name, requiredSize, Arrays.asList(topBucketsPerOrd.get(ordIdx)), metadata())
+            );
         }
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/prefix/IpPrefixAggregator.java

@@ -172,32 +172,32 @@ public final class IpPrefixAggregator extends BucketsAggregator {
             }
 
             try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(totalOrdsToCollect)) {
-                int b = 0;
+                int[] b = new int[] { 0 };
                 for (long i = 0; i < owningBucketOrds.size(); i++) {
                     BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(i));
                     while (ordsEnum.next()) {
-                        bucketOrdsToCollect.set(b++, ordsEnum.ord());
+                        bucketOrdsToCollect.set(b[0]++, ordsEnum.ord());
                     }
                 }
 
                 var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
-                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-                b = 0;
-                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                b[0] = 0;
+                return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                     List<InternalIpPrefix.Bucket> buckets = new ArrayList<>(bucketsInOrd.get(ordIdx));
                     BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
                     while (ordsEnum.next()) {
                         long ordinal = ordsEnum.ord();
-                        if (bucketOrdsToCollect.get(b) != ordinal) {
+                        if (bucketOrdsToCollect.get(b[0]) != ordinal) {
                             throw AggregationErrors.iterationOrderChangedWithoutMutating(
                                 bucketOrds.toString(),
                                 ordinal,
-                                bucketOrdsToCollect.get(b)
+                                bucketOrdsToCollect.get(b[0])
                             );
                         }
                         BytesRef ipAddress = new BytesRef();
                         ordsEnum.readValue(ipAddress);
                         long docCount = bucketDocCount(ordinal);
+                        checkRealMemoryCBForInternalBucket();
                         buckets.add(
                             new InternalIpPrefix.Bucket(
                                 config.format(),
@@ -207,16 +207,15 @@ public final class IpPrefixAggregator extends BucketsAggregator {
                                 ipPrefix.prefixLength,
                                 ipPrefix.appendPrefixLength,
                                 docCount,
-                                subAggregationResults.apply(b++)
+                                subAggregationResults.apply(b[0]++)
                             )
                         );
 
                         // NOTE: the aggregator is expected to return sorted results
                         CollectionUtil.introSort(buckets, BucketOrder.key(true).comparator());
                     }
-                    results[ordIdx] = new InternalIpPrefix(name, config.format(), keyed, minDocCount, buckets, metadata());
-                }
-                return results;
+                    return new InternalIpPrefix(name, config.format(), keyed, minDocCount, buckets, metadata());
+                });
             }
         }
     }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -696,11 +696,10 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
 
             if (valueCount == 0) { // no context in this reader
-                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-                    results[ordIdx] = buildNoValuesResult(owningBucketOrds.get(ordIdx));
-                }
-                return results;
+                return GlobalOrdinalsStringTermsAggregator.this.buildAggregations(
+                    Math.toIntExact(owningBucketOrds.size()),
+                    ordIdx -> buildNoValuesResult(owningBucketOrds.get(ordIdx))
+                );
             }
             try (
                 LongArray otherDocCount = bigArrays().newLongArray(owningBucketOrds.size(), true);
@@ -727,6 +726,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
                                 otherDocCount.increment(finalOrdIdx, docCount);
                                 if (docCount >= bucketCountThresholds.getShardMinDocCount()) {
                                     if (spare == null) {
+                                        checkRealMemoryCBForInternalBucket();
                                         spare = buildEmptyTemporaryBucket();
                                     }
                                     updater.updateBucket(spare, globalOrd, bucketOrd, docCount);
@@ -738,6 +738,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
                         // Get the top buckets
                         topBucketsPreOrd.set(ordIdx, buildBuckets((int) ordered.size()));
                         for (int i = (int) ordered.size() - 1; i >= 0; --i) {
+                            checkRealMemoryCBForInternalBucket();
                             B bucket = convertTempBucketToRealBucket(ordered.pop(), lookupGlobalOrd);
                             topBucketsPreOrd.get(ordIdx)[i] = bucket;
                             otherDocCount.increment(ordIdx, -bucket.getDocCount());
@@ -747,11 +748,10 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
 
                 buildSubAggs(topBucketsPreOrd);
 
-                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(topBucketsPreOrd.size())];
-                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-                    results[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCount.get(ordIdx), topBucketsPreOrd.get(ordIdx));
-                }
-                return results;
+                return GlobalOrdinalsStringTermsAggregator.this.buildAggregations(
+                    Math.toIntExact(owningBucketOrds.size()),
+                    ordIdx -> buildResult(owningBucketOrds.get(ordIdx), otherDocCount.get(ordIdx), topBucketsPreOrd.get(ordIdx))
+                );
             }
         }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java

@@ -142,6 +142,7 @@ public class LongRareTermsAggregator extends AbstractRareTermsAggregator {
                             long docCount = bucketDocCount(collectedBuckets.ord());
                             // if the key is below threshold, reinsert into the new ords
                             if (docCount <= maxDocCount) {
+                                checkRealMemoryCBForInternalBucket();
                                 LongRareTerms.Bucket bucket = new LongRareTerms.Bucket(collectedBuckets.value(), docCount, null, format);
                                 bucket.bucketOrd = offset + bucketsInThisOwningBucketToCollect.add(collectedBuckets.value());
                                 mergeMap.set(collectedBuckets.ord(), bucket.bucketOrd);
@@ -173,21 +174,12 @@ public class LongRareTermsAggregator extends AbstractRareTermsAggregator {
              * Now build the results!
              */
             buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+
+            return LongRareTermsAggregator.this.buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                 LongRareTerms.Bucket[] buckets = rarestPerOrd.get(ordIdx);
                 Arrays.sort(buckets, ORDER.comparator());
-                result[ordIdx] = new LongRareTerms(
-                    name,
-                    ORDER,
-                    metadata(),
-                    format,
-                    Arrays.asList(buckets),
-                    maxDocCount,
-                    filters.get(ordIdx)
-                );
-            }
-            return result;
+                return new LongRareTerms(name, ORDER, metadata(), format, Arrays.asList(buckets), maxDocCount, filters.get(ordIdx));
+            });
         }
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/MapStringTermsAggregator.java

@@ -304,6 +304,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
                                 continue;
                             }
                             if (spare == null) {
+                                checkRealMemoryCBForInternalBucket();
                                 spare = emptyBucketBuilder.get();
                             }
                             updateBucket(spare, ordsEnum, docCount);
@@ -320,11 +321,11 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
                 }
 
                 buildSubAggs(topBucketsPerOrd);
-                InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
-                for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
-                    result[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
-                }
-                return result;
+
+                return MapStringTermsAggregator.this.buildAggregations(
+                    Math.toIntExact(owningBucketOrds.size()),
+                    ordIdx -> buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx))
+                );
             }
         }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/NumericTermsAggregator.java

@@ -185,6 +185,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
                                 continue;
                             }
                             if (spare == null) {
+                                checkRealMemoryCBForInternalBucket();
                                 spare = emptyBucketBuilder.get();
                             }
                             updateBucket(spare, ordsEnum, docCount);
@@ -203,11 +204,10 @@ public final class NumericTermsAggregator extends TermsAggregator {
 
                 buildSubAggs(topBucketsPerOrd);
 
-                InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
-                for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
-                    result[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
-                }
-                return result;
+                return NumericTermsAggregator.this.buildAggregations(
+                    Math.toIntExact(owningBucketOrds.size()),
+                    ordIdx -> buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx))
+                );
             }
         }
 

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/StringRareTermsAggregator.java

@@ -145,6 +145,7 @@ public class StringRareTermsAggregator extends AbstractRareTermsAggregator {
                             long docCount = bucketDocCount(collectedBuckets.ord());
                             // if the key is below threshold, reinsert into the new ords
                             if (docCount <= maxDocCount) {
+                                checkRealMemoryCBForInternalBucket();
                                 StringRareTerms.Bucket bucket = new StringRareTerms.Bucket(
                                     BytesRef.deepCopyOf(scratch),
                                     docCount,
@@ -181,21 +182,12 @@ public class StringRareTermsAggregator extends AbstractRareTermsAggregator {
              * Now build the results!
              */
             buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+
+            return StringRareTermsAggregator.this.buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> {
                 StringRareTerms.Bucket[] buckets = rarestPerOrd.get(ordIdx);
                 Arrays.sort(buckets, ORDER.comparator());
-                result[ordIdx] = new StringRareTerms(
-                    name,
-                    ORDER,
-                    metadata(),
-                    format,
-                    Arrays.asList(buckets),
-                    maxDocCount,
-                    filters.get(ordIdx)
-                );
-            }
-            return result;
+                return new StringRareTerms(name, ORDER, metadata(), format, Arrays.asList(buckets), maxDocCount, filters.get(ordIdx));
+            });
         }
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/metrics/MetricsAggregator.java

@@ -38,10 +38,6 @@ public abstract class MetricsAggregator extends AggregatorBase {
 
     @Override
     public final InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
-        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-            results[ordIdx] = buildAggregation(owningBucketOrds.get(ordIdx));
-        }
-        return results;
+        return buildAggregations(Math.toIntExact(owningBucketOrds.size()), ordIdx -> buildAggregation(owningBucketOrds.get(ordIdx)));
     }
 }

x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/multiterms/MultiTermsAggregator.java

@@ -264,6 +264,7 @@ class MultiTermsAggregator extends DeferableBucketAggregator {
                             continue;
                         }
                         if (spare == null) {
+                            checkRealMemoryCBForInternalBucket();
                             spare = new InternalMultiTerms.Bucket(null, 0, null, showTermDocCountError, 0, formats, keyConverters);
                             spareKey = new BytesRef();
                         }
@@ -287,11 +288,10 @@ class MultiTermsAggregator extends DeferableBucketAggregator {
 
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
 
-            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
-            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
-                result[ordIdx] = buildResult(otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
-            }
-            return result;
+            return buildAggregations(
+                Math.toIntExact(owningBucketOrds.size()),
+                ordIdx -> buildResult(otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx))
+            );
         }
     }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java

@@ -121,21 +121,22 @@ public class CategorizeTextAggregator extends DeferableBucketAggregator {
                     continue;
                 }
                 int size = (int) Math.min(bucketOrds.bucketsInOrd(ordIdx), bucketCountThresholds.getShardSize());
+                checkRealMemoryCBForInternalBucket();
                 topBucketsPerOrd.set(ordIdx, categorizer.toOrderedBuckets(size));
             }
             buildSubAggsForAllBuckets(topBucketsPerOrd, Bucket::getBucketOrd, Bucket::setAggregations);
-            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(ordsToCollect.size())];
-            for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
-                results[ordIdx] = new InternalCategorizationAggregation(
+
+            return buildAggregations(
+                Math.toIntExact(ordsToCollect.size()),
+                ordIdx -> new InternalCategorizationAggregation(
                     name,
                     bucketCountThresholds.getRequiredSize(),
                     bucketCountThresholds.getMinDocCount(),
                     similarityThreshold,
                     metadata(),
                     Arrays.asList(topBucketsPerOrd.get(ordIdx))
-                );
-            }
-            return results;
+                )
+            );
         }
     }