Update breaker in tests

ChrisHegarty · Sep 22, 2023 · 61d3196 · 61d3196
1 parent 74d9b1d
commit 61d3196
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 108 deletions.
diff --git a/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/CountingCircuitBreaker.java b/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/CountingCircuitBreaker.java
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java
@@ -29,6 +29,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.MockPageCacheRecycler;
@@ -398,7 +399,8 @@ private BigArrays bigArrays() {
      * A {@link DriverContext} that won't throw {@link CircuitBreakingException}.
      */
     protected final DriverContext driverContext() {
-        return new DriverContext(bigArrays(), BlockFactory.getInstance(new CountingCircuitBreaker("esql-test"), bigArrays()));
+        var breaker = new MockBigArrays.LimitedBreaker("esql-test-breaker", ByteSizeValue.ofGb(1));
+        return new DriverContext(bigArrays(), BlockFactory.getInstance(breaker, bigArrays()));
     }
 
     public static void assertDriverContext(DriverContext driverContext) {

diff --git a/...ugin/esql/compute/src/test/java/org/elasticsearch/compute/TestBlockFactoryParameters.java b/...ugin/esql/compute/src/test/java/org/elasticsearch/compute/TestBlockFactoryParameters.java
@@ -8,6 +8,7 @@
 package org.elasticsearch.compute;
 
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
@@ -23,8 +24,7 @@ public class TestBlockFactoryParameters implements BlockFactoryParameters {
     final BigArrays bigArrays;
 
     public TestBlockFactoryParameters() {
-        breaker = new CountingCircuitBreaker("ESQL-test-breaker");
-
+        breaker = new MockBigArrays.LimitedBreaker("esql-test-breaker", ByteSizeValue.ofGb(1));
         var breakerService = mock(CircuitBreakerService.class);
         when(breakerService.getBreaker(CircuitBreaker.REQUEST)).thenReturn(breaker);
         bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, breakerService);

diff --git a/...c/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java b/...c/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java
@@ -10,14 +10,11 @@
 import com.carrotsearch.randomizedtesting.annotations.Name;
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
-import org.elasticsearch.compute.CountingCircuitBreaker;
 import org.elasticsearch.compute.data.BasicBlockTests;
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.operator.HashAggregationOperator;
 import org.elasticsearch.compute.operator.MultivalueDedupeTests;
@@ -40,9 +37,6 @@
 //@TestLogging(value = "org.elasticsearch.compute:TRACE", reason = "debug")
 public class BlockHashRandomizedTests extends ESTestCase {
 
-    static final CircuitBreaker breaker = new CountingCircuitBreaker("ESQL-test-breaker");
-    static final BlockFactory blockFactory = BlockFactory.getInstance(breaker, BigArrays.NON_RECYCLING_INSTANCE);
-
     @ParametersFactory
     public static List<Object[]> params() {
         List<Object[]> params = new ArrayList<>();

diff --git a/...compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java b/...compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java
@@ -12,10 +12,10 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.inject.name.Named;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
-import org.elasticsearch.compute.CountingCircuitBreaker;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
@@ -49,7 +49,7 @@
 
 public class BlockHashTests extends ESTestCase {
 
-    static final CircuitBreaker breaker = new CountingCircuitBreaker("ESQL-test-breaker");
+    static final CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("esql-test-breaker", ByteSizeValue.ofGb(1));
     static final BlockFactory blockFactory = BlockFactory.getInstance(breaker, BigArrays.NON_RECYCLING_INSTANCE);
 
     @ParametersFactory

diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicBlockTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicBlockTests.java
@@ -10,11 +10,11 @@
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
-import org.elasticsearch.compute.CountingCircuitBreaker;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
@@ -41,14 +41,9 @@
 
 public class BasicBlockTests extends ESTestCase {
 
-    final CircuitBreaker breaker = new CountingCircuitBreaker("ESQL-test-breaker");
-    final BlockFactory blockFactory = BlockFactory.getInstance(breaker, bigArrays());
-
-    BigArrays bigArrays() {
-        var breakerService = mock(CircuitBreakerService.class);
-        when(breakerService.getBreaker(CircuitBreaker.REQUEST)).thenReturn(breaker);
-        return new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, breakerService);
-    }
+    final CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("esql-test-breaker", ByteSizeValue.ofGb(1));
+    final BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, mockBreakerService(breaker));
+    final BlockFactory blockFactory = BlockFactory.getInstance(breaker, bigArrays);
 
     @Before
     @After
@@ -60,11 +55,7 @@ public void testEmpty() {
         testEmpty(blockFactory);
     }
 
-    public void testEmptyNonBreakingFactory() {
-        testEmpty(BlockFactory.getNonBreakingInstance());
-    }
-
-    public void testEmpty(BlockFactory bf) {
+    void testEmpty(BlockFactory bf) {
         assertZeroPositionsAndRelease(bf.newIntArrayBlock(new int[] {}, 0, new int[] {}, new BitSet(), randomOrdering()));
         assertZeroPositionsAndRelease(IntBlock.newBlockBuilder(0, bf).build());
         assertZeroPositionsAndRelease(bf.newIntArrayVector(new int[] {}, 0));
@@ -252,7 +243,7 @@ public void testLongBlock() {
             if (positionCount > 1) {
                 assertNullValues(
                     positionCount,
-                    size -> LongBlock.newBlockBuilder(size),
+                    size -> LongBlock.newBlockBuilder(size, blockFactory),
                     (bb, value) -> bb.appendLong(value),
                     position -> (long) position,
                     LongBlock.Builder::build,
@@ -323,7 +314,7 @@ public void testDoubleBlock() {
             if (positionCount > 1) {
                 assertNullValues(
                     positionCount,
-                    size -> DoubleBlock.newBlockBuilder(size),
+                    size -> DoubleBlock.newBlockBuilder(size, blockFactory),
                     (bb, value) -> bb.appendDouble(value),
                     position -> (double) position,
                     DoubleBlock.Builder::build,
@@ -400,7 +391,7 @@ public void testBytesRefBlock() {
         if (positionCount > 1) {
             assertNullValues(
                 positionCount,
-                size -> BytesRefBlock.newBlockBuilder(size),
+                size -> BytesRefBlock.newBlockBuilder(size, blockFactory),
                 (bb, value) -> bb.appendBytesRef(value),
                 position -> values[position],
                 BytesRefBlock.Builder::build,
@@ -510,7 +501,7 @@ public void testBooleanBlock() {
             if (positionCount > 1) {
                 assertNullValues(
                     positionCount,
-                    BooleanBlock::newBlockBuilder,
+                    size -> BooleanBlock.newBlockBuilder(size, blockFactory),
                     (bb, value) -> bb.appendBoolean(value),
                     position -> position % 10 == 0,
                     BooleanBlock.Builder::build,
@@ -912,6 +903,7 @@ <T extends Releasable & Accountable> void releaseAndAssertBreaker(T data) {
     }
 
     static <T extends Releasable & Accountable> void releaseAndAssertBreaker(T data, CircuitBreaker breaker) {
+        assertThat(breaker.getUsed(), greaterThan(0L));
         Releasables.closeExpectNoException(data);
         assertThat(breaker.getUsed(), is(0L));
     }
@@ -923,4 +915,11 @@ static int randomPosition(int positionCount) {
     static Block.MvOrdering randomOrdering() {
         return randomFrom(Block.MvOrdering.values());
     }
+
+    // A breaker service that always returns the given breaker for getBreaker(CircuitBreaker.REQUEST)
+    static CircuitBreakerService mockBreakerService(CircuitBreaker breaker) {
+        CircuitBreakerService breakerService = mock(CircuitBreakerService.class);
+        when(breakerService.getBreaker(CircuitBreaker.REQUEST)).thenReturn(breaker);
+        return breakerService;
+    }
 }
diff --git a/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java b/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java
@@ -13,11 +13,11 @@
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
-import org.elasticsearch.compute.CountingCircuitBreaker;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
@@ -36,24 +36,17 @@
 
 // BlockFactory is used and effectively tested in many other places, but this class contains tests
 // more specific to the factory implementation itself (and not necessarily tested elsewhere).
-// @com.carrotsearch.randomizedtesting.annotations.Repeat(iterations = 1000)
 public class BlockFactoryTests extends ESTestCase {
 
     final CircuitBreaker breaker;
     final BigArrays bigArrays;
     final BlockFactory blockFactory;
 
-    static BigArrays bigArrays(CircuitBreaker breaker) {
-        var breakerService = mock(CircuitBreakerService.class);
-        when(breakerService.getBreaker(CircuitBreaker.REQUEST)).thenReturn(breaker);
-        return new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, breakerService);
-    }
-
     @ParametersFactory
     public static List<Object[]> params() {
         List<Supplier<BlockFactory>> l = List.of(() -> {
-            CircuitBreaker breaker = new CountingCircuitBreaker("ESQL-test-breaker");
-            BigArrays bigArrays = bigArrays(breaker);
+            CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("esql-test-breaker", ByteSizeValue.ofGb(1));
+            BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, mockBreakerService(breaker));
             return BlockFactory.getInstance(breaker, bigArrays);
         }, BlockFactory::getGlobalInstance);
         return l.stream().map(s -> new Object[] { s }).toList();
@@ -557,7 +550,15 @@ static Block.MvOrdering randomOrdering() {
     }
 
     <T extends Releasable & Accountable> void releaseAndAssertBreaker(T data) {
+        assertThat(breaker.getUsed(), greaterThan(0L));
         Releasables.closeExpectNoException(data);
         assertThat(breaker.getUsed(), is(0L));
     }
+
+    // A breaker service that always returns the given breaker for getBreaker(CircuitBreaker.REQUEST)
+    static CircuitBreakerService mockBreakerService(CircuitBreaker breaker) {
+        CircuitBreakerService breakerService = mock(CircuitBreakerService.class);
+        when(breakerService.getBreaker(CircuitBreaker.REQUEST)).thenReturn(breaker);
+        return breakerService;
+    }
 }