From a021210335121315dc68ff6ed80e6a4aea2b1817 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Thu, 26 Oct 2023 10:59:30 -0400 Subject: [PATCH] Default to tracked block factory --- .../AggregatorFunctionTestCase.java | 6 ------ .../GroupingAggregatorFunctionTestCase.java | 6 ------ .../lucene/LuceneSourceOperatorTests.java | 5 +++++ .../lucene/LuceneTopNSourceOperatorTests.java | 6 ++++++ .../ValuesSourceReaderOperatorTests.java | 5 ----- .../compute/operator/AnyOperatorTestCase.java | 19 ++++++++----------- .../compute/operator/EvalOperatorTests.java | 5 ----- .../compute/operator/FilterOperatorTests.java | 5 ----- .../HashAggregationOperatorTests.java | 6 ------ .../compute/operator/LimitOperatorTests.java | 5 ----- .../operator/MvExpandOperatorTests.java | 5 ----- .../operator/ProjectOperatorTests.java | 5 ----- .../operator/StringExtractOperatorTests.java | 5 +++++ .../operator/topn/TopNOperatorTests.java | 5 ----- 14 files changed, 24 insertions(+), 64 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java index 83f4a6895b154..846edf9fdedad 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java @@ -43,12 +43,6 @@ import static org.hamcrest.Matchers.hasSize; public abstract class AggregatorFunctionTestCase extends ForkingOperatorTestCase { - - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } - protected abstract AggregatorFunctionSupplier aggregatorFunction(BigArrays bigArrays, List inputChannels); protected final int aggregatorIntermediateBlockCount() { diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java index ba930e943e79a..753b5878de2ae 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java @@ -48,12 +48,6 @@ import static org.hamcrest.Matchers.hasSize; public abstract class GroupingAggregatorFunctionTestCase extends ForkingOperatorTestCase { - - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } - protected abstract AggregatorFunctionSupplier aggregatorFunction(BigArrays bigArrays, List inputChannels); protected final int aggregatorIntermediateBlockCount() { diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java index 131082859bf4c..12c6d19ce88a4 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java @@ -186,4 +186,9 @@ public static SearchContext mockSearchContext(IndexReader reader) { throw new UncheckedIOException(e); } } + + @Override + protected DriverContext driverContext() { + return nonBreakingDriverContext(); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java index f0eb49c233e7a..50c133e484472 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java @@ -171,4 +171,10 @@ private void testSimple(int size, int limit) { int pages = (int) Math.ceil((float) Math.min(size, limit) / factory.maxPageSize()); assertThat(results, hasSize(pages)); } + + @Override + protected DriverContext driverContext() { + return nonBreakingDriverContext(); + } + } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java index 622ef6d610b7d..269a478560bac 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java @@ -384,9 +384,4 @@ public void testValuesSourceReaderOperatorWithNulls() throws IOException { } assertDriverContext(driverContext); } - - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java index 00c054ef0031a..f1a9c144d9165 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java @@ -100,17 +100,7 @@ protected final BigArrays nonBreakingBigArrays() { /** * A {@link DriverContext} with a nonBreakingBigArrays. */ - protected DriverContext driverContext() { // TODO make this final and return a breaking block factory - return new DriverContext(nonBreakingBigArrays(), BlockFactory.getNonBreakingInstance()); - } - - private final List breakers = new ArrayList<>(); - private final List blockFactories = new ArrayList<>(); - - /** - * A {@link DriverContext} with a breaking {@link BigArrays} and {@link BlockFactory}. - */ - protected DriverContext breakingDriverContext() { // TODO move this to driverContext once everyone supports breaking + protected DriverContext driverContext() { // TODO make this final once all operators support memory tracking BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, ByteSizeValue.ofGb(1)).withCircuitBreaking(); CircuitBreaker breaker = bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST); breakers.add(breaker); @@ -119,6 +109,13 @@ protected DriverContext breakingDriverContext() { // TODO move this to driverCon return new DriverContext(bigArrays, factory); } + protected final DriverContext nonBreakingDriverContext() { // TODO drop this once the driverContext method isn't overrideable + return new DriverContext(nonBreakingBigArrays(), BlockFactory.getNonBreakingInstance()); + } + + private final List breakers = new ArrayList<>(); + private final List blockFactories = new ArrayList<>(); + @After public void allBreakersEmpty() throws Exception { // first check that all big arrays are released, which can affect breakers diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java index 95a5647717851..e7f5db7579869 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java @@ -118,9 +118,4 @@ public void testReadFromBlock() { protected ByteSizeValue smallEnoughToCircuitBreak() { return ByteSizeValue.ofBytes(between(1, 8000)); } - - @Override - protected DriverContext driverContext() { // TODO remove this when the parent uses a breaking block factory - return breakingDriverContext(); - } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java index fa4c7bea7c9cc..e16f643e1ca4d 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java @@ -119,9 +119,4 @@ public void testReadFromBlock() { protected ByteSizeValue smallEnoughToCircuitBreak() { return ByteSizeValue.ofBytes(between(1, 600)); } - - @Override - protected DriverContext driverContext() { // TODO remove this when the parent uses a breaking block factory - return breakingDriverContext(); - } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java index afa307c494431..b1ef784ca339c 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java @@ -31,12 +31,6 @@ import static org.hamcrest.Matchers.hasSize; public class HashAggregationOperatorTests extends ForkingOperatorTestCase { - - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } - @Override protected SourceOperator simpleInput(BlockFactory blockFactory, int size) { long max = randomLongBetween(1, Long.MAX_VALUE / size); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java index 76f99389a697b..8c85f5927196f 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java @@ -23,11 +23,6 @@ import static org.hamcrest.Matchers.sameInstance; public class LimitOperatorTests extends OperatorTestCase { - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } - @Override protected LimitOperator.Factory simple(BigArrays bigArrays) { return new LimitOperator.Factory(100); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java index a105818a2bf03..3572dc620287d 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java @@ -257,9 +257,4 @@ protected Page createPage(int positionOffset, int length) { List results = drive(new MvExpandOperator(0, randomIntBetween(1, 1000)), input.iterator(), context); assertSimpleOutput(origInput, results); } - - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java index 1acdbc4895c94..1aff6be7594aa 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java @@ -28,11 +28,6 @@ import static org.mockito.Mockito.when; public class ProjectOperatorTests extends OperatorTestCase { - @Override - protected DriverContext driverContext() { - return breakingDriverContext(); - } - public void testProjectionOnEmptyPage() { var page = new Page(0); var projection = new ProjectOperator(randomProjection(10)); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java index 5dff00ec930c4..55470d1dcae12 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java @@ -126,4 +126,9 @@ public void close() {} assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo4")); assertThat(brb.getBytesRef(idx + 2, spare).utf8ToString(), equalTo("foo5")); } + + @Override + protected DriverContext driverContext() { + return nonBreakingDriverContext(); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java index 049bc449069b8..6c5bab9b8f784 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java @@ -1359,11 +1359,6 @@ public void testCloseWithoutCompleting() { } } - @Override - protected DriverContext driverContext() { // TODO remove this when the parent uses a breaking block factory - return breakingDriverContext(); - } - @SuppressWarnings({ "unchecked", "rawtypes" }) private static void readAsRows(List>> values, Page page) { if (page.getBlockCount() == 0) {