From 88785f14df3ea2927bffc8ed55251c32c3371f94 Mon Sep 17 00:00:00 2001 From: bowenlan-amzn Date: Fri, 19 Jul 2024 19:13:37 -0700 Subject: [PATCH] Use circuit breaker in InternalHistogram when adding empty buckets (#14754) (#14844) Signed-off-by: kkewwei --- CHANGELOG.md | 1 + .../bucket/histogram/InternalHistogram.java | 6 ++- .../histogram/InternalHistogramTests.java | 43 +++++++++++++++++++ 3 files changed, 49 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 988b9917a5471..70fedcd8fdda1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Update help output for _cat ([#14722](https://github.com/opensearch-project/OpenSearch/pull/14722)) - Fix bulk upsert ignores the default_pipeline and final_pipeline when auto-created index matches the index template ([#12891](https://github.com/opensearch-project/OpenSearch/pull/12891)) - Fix NPE in ReplicaShardAllocator ([#14385](https://github.com/opensearch-project/OpenSearch/pull/14385)) +- Use circuit breaker in InternalHistogram when adding empty buckets ([#14754](https://github.com/opensearch-project/OpenSearch/pull/14754)) ### Security diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogram.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogram.java index a27c689127ac9..a988b911de5a3 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogram.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogram.java @@ -395,6 +395,7 @@ private void addEmptyBuckets(List list, ReduceContext reduceContext) { // fill with empty buckets for (double key = round(emptyBucketInfo.minBound); key <= emptyBucketInfo.maxBound; key = nextKey(key)) { iter.add(new Bucket(key, 0, keyed, format, reducedEmptySubAggs)); + reduceContext.consumeBucketsAndMaybeBreak(0); } } else { Bucket first = list.get(iter.nextIndex()); @@ -402,11 +403,12 @@ private void addEmptyBuckets(List list, ReduceContext reduceContext) { // fill with empty buckets until the first key for (double key = round(emptyBucketInfo.minBound); key < first.key; key = nextKey(key)) { iter.add(new Bucket(key, 0, keyed, format, reducedEmptySubAggs)); + reduceContext.consumeBucketsAndMaybeBreak(0); } } // now adding the empty buckets within the actual data, - // e.g. if the data series is [1,2,3,7] there're 3 empty buckets that will be created for 4,5,6 + // e.g. if the data series is [1,2,3,7] there are 3 empty buckets that will be created for 4,5,6 Bucket lastBucket = null; do { Bucket nextBucket = list.get(iter.nextIndex()); @@ -414,6 +416,7 @@ private void addEmptyBuckets(List list, ReduceContext reduceContext) { double key = nextKey(lastBucket.key); while (key < nextBucket.key) { iter.add(new Bucket(key, 0, keyed, format, reducedEmptySubAggs)); + reduceContext.consumeBucketsAndMaybeBreak(0); key = nextKey(key); } assert key == nextBucket.key || Double.isNaN(nextBucket.key) : "key: " + key + ", nextBucket.key: " + nextBucket.key; @@ -424,6 +427,7 @@ private void addEmptyBuckets(List list, ReduceContext reduceContext) { // finally, adding the empty buckets *after* the actual data (based on the extended_bounds.max requested by the user) for (double key = nextKey(lastBucket.key); key <= emptyBucketInfo.maxBound; key = nextKey(key)) { iter.add(new Bucket(key, 0, keyed, format, reducedEmptySubAggs)); + reduceContext.consumeBucketsAndMaybeBreak(0); } } } diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogramTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogramTests.java index 288b22ccfcc92..98c6ac2b3de45 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogramTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/histogram/InternalHistogramTests.java @@ -33,10 +33,15 @@ package org.opensearch.search.aggregations.bucket.histogram; import org.apache.lucene.tests.util.TestUtil; +import org.opensearch.core.common.breaker.CircuitBreaker; +import org.opensearch.core.common.breaker.CircuitBreakingException; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.BucketOrder; +import org.opensearch.search.aggregations.InternalAggregation; import org.opensearch.search.aggregations.InternalAggregations; +import org.opensearch.search.aggregations.MultiBucketConsumerService; import org.opensearch.search.aggregations.ParsedMultiBucketAggregation; +import org.opensearch.search.aggregations.pipeline.PipelineAggregator; import org.opensearch.test.InternalAggregationTestCase; import org.opensearch.test.InternalMultiBucketAggregationTestCase; @@ -47,6 +52,8 @@ import java.util.Map; import java.util.TreeMap; +import org.mockito.Mockito; + public class InternalHistogramTests extends InternalMultiBucketAggregationTestCase { private boolean keyed; @@ -123,6 +130,42 @@ public void testHandlesNaN() { ); } + public void testCircuitBreakerWhenAddEmptyBuckets() { + String name = randomAlphaOfLength(5); + double interval = 1; + double lowerBound = 1; + double upperBound = 1026; + List bucket1 = List.of( + new InternalHistogram.Bucket(lowerBound, 1, false, format, InternalAggregations.EMPTY) + ); + List bucket2 = List.of( + new InternalHistogram.Bucket(upperBound, 1, false, format, InternalAggregations.EMPTY) + ); + BucketOrder order = BucketOrder.key(true); + InternalHistogram.EmptyBucketInfo emptyBucketInfo = new InternalHistogram.EmptyBucketInfo( + interval, + 0, + lowerBound, + upperBound, + InternalAggregations.EMPTY + ); + InternalHistogram histogram1 = new InternalHistogram(name, bucket1, order, 0, emptyBucketInfo, format, false, null); + InternalHistogram histogram2 = new InternalHistogram(name, bucket2, order, 0, emptyBucketInfo, format, false, null); + + CircuitBreaker breaker = Mockito.mock(CircuitBreaker.class); + Mockito.when(breaker.addEstimateBytesAndMaybeBreak(0, "allocated_buckets")).thenThrow(CircuitBreakingException.class); + + MultiBucketConsumerService.MultiBucketConsumer bucketConsumer = new MultiBucketConsumerService.MultiBucketConsumer(0, breaker); + InternalAggregation.ReduceContext reduceContext = InternalAggregation.ReduceContext.forFinalReduction( + null, + null, + bucketConsumer, + PipelineAggregator.PipelineTree.EMPTY + ); + expectThrows(CircuitBreakingException.class, () -> histogram1.reduce(List.of(histogram1, histogram2), reduceContext)); + Mockito.verify(breaker, Mockito.times(1)).addEstimateBytesAndMaybeBreak(0, "allocated_buckets"); + } + @Override protected void assertReduced(InternalHistogram reduced, List inputs) { TreeMap expectedCounts = new TreeMap<>();