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

checks periodically the real memory circuit breaker when allocating objects.

Author: iverase

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

@@ -188,17 +188,16 @@ public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
             }
         }
         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 InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                     // 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 InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                             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 InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                 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 InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
             }
             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.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 @@ protected AggregatorBase(
         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 @@ protected final InternalAggregations buildEmptySubAggregations() {
         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 final LeafBucketCollector getLeafCollector(AggregationExecutionContext ag
 
     @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.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 BucketsAggregator(
         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 final void collectBucket(LeafBucketCollector subCollector, int doc, long
         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 @@ protected final IntFunction<InternalAggregations> buildSubAggsForBuckets(LongArr
         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 @@ protected final <B> InternalAggregation[] buildAggregationsForFixedBucketCount(
         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 @@ protected final InternalAggregation[] buildAggregationsForSingleBucket(
          * 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 @@ protected final <B> InternalAggregation[] buildAggregationsForVariableBuckets(
                 );
             }
             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 @@ protected void preGetSubLeafCollectors(LeafReaderContext ctx) throws IOException
         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 @@ public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                         long docCount = bucketDocCount(ordsEnum.ord());
                         otherDocCounts.increment(ordIdx, docCount);
                         if (spare == null) {
+                            checkRealMemoryCBForInternalBucket();
                             spare = emptyBucketBuilder.get();
                         }
                         ordsEnum.readValue(spare.getTermBytes());
@@ -158,16 +159,16 @@ public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
             }
 
             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);
                     Arrays.sort(topBucketsPerOrd.get(ordIdx), reduceOrder.comparator());
                 } else {
                     reduceOrder = order;
                 }
-                result[ordIdx] = new StringTerms(
+                return new StringTerms(
                     name,
                     reduceOrder,
                     order,
@@ -181,8 +182,7 @@ public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                     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 InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
                     while (ordsEnum.next()) {
                         if (spare == null) {
+                            checkRealMemoryCBForInternalBucket();
                             spare = newEmptyBucket();
                         }
 
@@ -162,11 +163,10 @@ public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throw
                 }
             }
             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())
+            );
         }
     }