Default to tracked block factory

nik9000 · Oct 26, 2023 · a021210 · a021210
1 parent a3efe17
commit a021210
Show file tree

Hide file tree

Showing 14 changed files with 24 additions and 64 deletions.
diff --git a/...mpute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java b/...mpute/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<Integer> inputChannels);
 
     protected final int aggregatorIntermediateBlockCount() {

diff --git a/...c/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java b/...c/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<Integer> inputChannels);
 
     protected final int aggregatorIntermediateBlockCount() {

diff --git a/...sql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/...sql/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/...compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/...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/...mpute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java b/...mpute/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/...in/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java b/...in/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<CircuitBreaker> breakers = new ArrayList<>();
-    private final List<BlockFactory> 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<CircuitBreaker> breakers = new ArrayList<>();
+    private final List<BlockFactory> blockFactories = new ArrayList<>();
+
     @After
     public void allBreakersEmpty() throws Exception {
         // first check that all big arrays are released, which can affect breakers

diff --git a/...ugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java b/...ugin/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/...in/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java b/...in/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/...ompute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java b/...ompute/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/...gin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java b/...gin/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/.../esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java b/.../esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java
@@ -257,9 +257,4 @@ protected Page createPage(int positionOffset, int length) {
         List<Page> results = drive(new MvExpandOperator(0, randomIntBetween(1, 1000)), input.iterator(), context);
         assertSimpleOutput(origInput, results);
     }
-
-    @Override
-    protected DriverContext driverContext() {
-        return breakingDriverContext();
-    }
 }
diff --git a/...n/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java b/...n/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/.../compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java b/.../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/...esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java b/...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<List<List<Object>>> values, Page page) {
         if (page.getBlockCount() == 0) {