diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/CountingCircuitBreaker.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/CountingCircuitBreaker.java deleted file mode 100644 index dc4f77070ffa0..0000000000000 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/CountingCircuitBreaker.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.compute; - -import org.elasticsearch.common.breaker.CircuitBreaker; - -public class CountingCircuitBreaker implements CircuitBreaker { - - private final String name; - - private long total; - - public CountingCircuitBreaker(String name) { - this.name = name; - } - - @Override - public void circuitBreak(String fieldName, long bytesNeeded) { - // noop - } - - @Override - public void addEstimateBytesAndMaybeBreak(long bytes, String label) { - total += bytes; - } - - @Override - public void addWithoutBreaking(long bytes) { - total += bytes; - } - - @Override - public long getUsed() { - return total; - } - - @Override - public long getLimit() { - return -1; - } - - @Override - public double getOverhead() { - return 0; - } - - @Override - public long getTrippedCount() { - return 0; - } - - @Override - public String getName() { - return this.name; - } - - @Override - public Durability getDurability() { - return Durability.PERMANENT; - } - - @Override - public void setLimitAndOverhead(long limit, double overhead) {} -} 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 index a182df8cb42ee..04a966b399870 100644 --- 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/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/TestBlockFactoryParameters.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/TestBlockFactoryParameters.java index a5957035aef79..8fa38b6864674 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/TestBlockFactoryParameters.java +++ b/x-pack/plugin/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/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java index af34d5733de1f..5775a983a2373 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java +++ b/x-pack/plugin/esql/compute/src/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 params() { List params = new ArrayList<>(); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java index 6e7aa383b381b..620bb5ab5319a 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java +++ b/x-pack/plugin/esql/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 index 2733f1a43eba6..df16f0036c767 100644 --- 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 @@ void releaseAndAssertBreaker(T data) { } static 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/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java index 87d6f595328e4..831be03cb0c81 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java +++ b/x-pack/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 params() { List> 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() { } 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; + } }