ESQL: Fix memory tests (#104082)

This improves the memory tests of ESQL by removing the hard memory
limits, instead calculating "how much memory is too much" on the
fly. We do a binary search on the operation itself, looking for
the first value that'll throw.

Closes #103789

test/framework/src/main/java/org/elasticsearch/test/BreakerTestUtil.java

@@ -0,0 +1,85 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.test;
+
+import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.core.CheckedConsumer;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+
+public class BreakerTestUtil {
+    private static final Logger logger = LogManager.getLogger(BreakerTestUtil.class);
+
+    /**
+     * Performs a binary search between 0 and {@code tooBigToBreak} bytes for the largest memory size
+     * that'll cause the closure parameter to throw a {@link CircuitBreakingException}.
+     */
+    public static <E extends Exception> ByteSizeValue findBreakerLimit(ByteSizeValue tooBigToBreak, CheckedConsumer<ByteSizeValue, E> c)
+        throws E {
+
+        // Validate arguments: we don't throw for tooBigToBreak and we *do* throw for 0.
+        try {
+            c.accept(tooBigToBreak);
+        } catch (CircuitBreakingException e) {
+            throw new IllegalArgumentException("expected runnable *not* to break under tooBigToBreak", e);
+        }
+        try {
+            c.accept(ByteSizeValue.ofBytes(0));
+            throw new IllegalArgumentException("expected runnable to break under a limit of 0 bytes");
+        } catch (CircuitBreakingException e) {
+            // desired
+        }
+
+        // Perform the actual binary search
+        long l = findBreakerLimit(0, tooBigToBreak.getBytes(), c);
+
+        // Validate results: we *do* throw for limit, we don't throw for limit + 1
+        ByteSizeValue limit = ByteSizeValue.ofBytes(l);
+        ByteSizeValue onePastLimit = ByteSizeValue.ofBytes(l + 1);
+        try {
+            c.accept(limit);
+            throw new IllegalArgumentException("expected runnable to break under a limit of " + limit + " bytes");
+        } catch (CircuitBreakingException e) {
+            // desired
+        }
+        try {
+            c.accept(onePastLimit);
+        } catch (CircuitBreakingException e) {
+            throw new IllegalArgumentException("expected runnable to break under a limit of " + onePastLimit + " bytes");
+        }
+        return limit;
+    }
+
+    /**
+     * A binary search of memory limits, looking for the lowest limit that'll break.
+     */
+    private static <E extends Exception> long findBreakerLimit(long min, long max, CheckedConsumer<ByteSizeValue, E> c) throws E {
+        // max is an amount of memory that doesn't break
+        // min is an amount of memory that *does* break
+        while (max - min > 1) {
+            assert max > min;
+            long diff = max - min;
+            logger.info(
+                "Between {} and {}. {} bytes remaining.",
+                ByteSizeValue.ofBytes(min),
+                ByteSizeValue.ofBytes(max),
+                ByteSizeValue.ofBytes(diff)
+            );
+            long mid = min + diff / 2;
+            try {
+                c.accept(ByteSizeValue.ofBytes(mid));
+                max = mid;
+            } catch (CircuitBreakingException e) {
+                min = mid;
+            }
+        }
+        return min;
+    }
+}

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockTestUtils;
@@ -82,12 +81,6 @@ public abstract class AggregatorFunctionTestCase extends ForkingOperatorTestCase
         assertSimpleOutput(input.stream().map(p -> p.<Block>getBlock(0)).toList(), result);
     }
 
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        // This is a super conservative limit that should cause all aggs to break
-        return ByteSizeValue.ofBytes(20);
-    }
-
     public final void testIgnoresNulls() {
         int end = between(1_000, 100_000);
         List<Page> results = new ArrayList<>();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.compute.aggregation;
 
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
@@ -144,11 +143,6 @@ public abstract class GroupingAggregatorFunctionTestCase extends ForkingOperator
         }
     }
 
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofBytes(100);
-    }
-
     public final void testNullGroupsAndValues() {
         DriverContext driverContext = driverContext();
         BlockFactory blockFactory = driverContext.blockFactory();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java

@@ -395,7 +395,7 @@ public class ValuesSourceReaderOperatorTests extends OperatorTestCase {
     }
 
     @Override
-    protected ByteSizeValue memoryLimitForSimple() {
+    protected ByteSizeValue enoughMemoryForSimple() {
         assumeFalse("strange exception in the test, fix soon", true);
         return ByteSizeValue.ofKb(1);
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AggregationOperatorTests.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.aggregation.MaxLongAggregatorFunction;
 import org.elasticsearch.compute.aggregation.MaxLongAggregatorFunctionSupplier;
@@ -80,9 +79,4 @@ public class AggregationOperatorTests extends ForkingOperatorTestCase {
         sum.assertSimpleOutput(input.stream().map(p -> p.<Block>getBlock(0)).toList(), sums);
         max.assertSimpleOutput(input.stream().map(p -> p.<Block>getBlock(0)).toList(), maxs);
     }
-
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofBytes(50);
-    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java

@@ -9,7 +9,6 @@ package org.elasticsearch.compute.operator;
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.lucene.BytesRefs;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
@@ -95,11 +94,6 @@ public class ColumnExtractOperatorTests extends OperatorTestCase {
         }
     }
 
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofKb(15);
-    }
-
     public void testAllNullValues() {
         DriverContext driverContext = driverContext();
         BytesRef scratch = new BytesRef();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/EvalOperatorTests.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
@@ -114,9 +113,4 @@ public class EvalOperatorTests extends OperatorTestCase {
         results.forEach(Page::releaseBlocks);
         assertThat(context.breaker().getUsed(), equalTo(0L));
     }
-
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofKb(4);
-    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
@@ -113,9 +112,4 @@ public class FilterOperatorTests extends OperatorTestCase {
         results.forEach(Page::releaseBlocks);
         assertThat(context.breaker().getUsed(), equalTo(0L));
     }
-
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofKb(1);
-    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ForkingOperatorTestCase.java

@@ -11,7 +11,6 @@ import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.data.BlockTestUtils;
@@ -177,7 +176,6 @@ public abstract class ForkingOperatorTestCase extends OperatorTestCase {
     // @com.carrotsearch.randomizedtesting.annotations.Repeat(iterations = 100)
     public final void testManyInitialManyPartialFinalRunnerThrowing() throws Exception {
         DriverContext driverContext = driverContext();
-        BigArrays bigArrays = nonBreakingBigArrays();
         List<Page> input = CannedSourceOperator.collectPages(simpleInput(driverContext.blockFactory(), between(1_000, 100_000)));
         List<Page> results = new ArrayList<>();
 

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.aggregation.MaxLongAggregatorFunction;
 import org.elasticsearch.compute.aggregation.MaxLongAggregatorFunctionSupplier;
@@ -91,9 +90,4 @@ public class HashAggregationOperatorTests extends ForkingOperatorTestCase {
             max.assertSimpleGroup(input, maxs, i, group);
         }
     }
-
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofKb(1);
-    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorTests.java

@@ -50,7 +50,7 @@ public class LimitOperatorTests extends OperatorTestCase {
     }
 
     @Override
-    protected ByteSizeValue memoryLimitForSimple() {
+    protected ByteSizeValue enoughMemoryForSimple() {
         assumeFalse("doesn't allocate, just filters", true);
         return null;
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/MvExpandOperatorTests.java

@@ -200,7 +200,7 @@ public class MvExpandOperatorTests extends OperatorTestCase {
     }
 
     @Override
-    protected ByteSizeValue memoryLimitForSimple() {
+    protected ByteSizeValue enoughMemoryForSimple() {
         assumeFalse("doesn't throw in tests but probably should", true);
         return ByteSizeValue.ofBytes(1);
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OperatorTestCase.java

@@ -28,10 +28,10 @@ import org.elasticsearch.compute.data.TestBlockFactory;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.indices.CrankyCircuitBreakerService;
+import org.elasticsearch.test.BreakerTestUtil;
 import org.elasticsearch.threadpool.FixedExecutorBuilder;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
-import org.junit.AssumptionViolatedException;
 
 import java.util.ArrayList;
 import java.util.Iterator;
@@ -73,69 +73,64 @@ public abstract class OperatorTestCase extends AnyOperatorTestCase {
     }
 
     /**
-     * A {@link ByteSizeValue} that is small enough that running {@link #simple}
-     * on {@link #simpleInput} will exhaust the breaker and throw a
-     * {@link CircuitBreakingException}. We should make an effort to make this
-     * number as large as possible and still cause a break consistently so we get
-     * good test coverage. If the operator can't break then throw an
-     * {@link AssumptionViolatedException}.
+     * Enough memory for {@link #simple} not to throw a {@link CircuitBreakingException}.
+     * It's fine if this is <strong>much</strong> more memory than {@linkplain #simple} needs.
+     * When we want to make {@linkplain #simple} throw we'll find the precise amount of memory
+     * that'll make it throw with a binary search.
      */
-    protected abstract ByteSizeValue memoryLimitForSimple();
-
-    /**
-     * Run {@link #simple} with a circuit breaker limited to somewhere
-     * between 0 bytes and {@link #memoryLimitForSimple} and assert that
-     * it breaks in a sane way.
-     */
-    public final void testSimpleCircuitBreaking() {
-        testSimpleCircuitBreaking(ByteSizeValue.ofBytes(randomLongBetween(0, memoryLimitForSimple().getBytes())));
+    protected ByteSizeValue enoughMemoryForSimple() {
+        return ByteSizeValue.ofGb(1);
     }
 
     /**
-     * Run {@link #simple} with a circuit breaker configured limited to
-     * {@link #memoryLimitForSimple} and assert that it breaks in a sane way.
-     * <p>
-     *     This test helps to make sure that the limits set by
-     *     {@link #memoryLimitForSimple} aren't too large.
-     *     {@link #testSimpleCircuitBreaking}, with it's random configured
-     *     limit will use the actual maximum very rarely.
-     * </p>
+     * Run {@link #simple} with a circuit breaker many times, making sure all blocks
+     * are properly released. In particular, we perform a binary search to find the
+     * largest amount of memory that'll throw a {@link CircuitBreakingException} with
+     * starting bounds of {@code 0b} and {@link #enoughMemoryForSimple}. Then we pick
+     * a random amount of memory between {@code 0b} and the maximum and run that,
+     * asserting both that this throws a {@link CircuitBreakingException} and releases
+     * all pages.
      */
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/103789")
-    public final void testSimpleCircuitBreakingAtLimit() {
-        testSimpleCircuitBreaking(memoryLimitForSimple());
+    public final void testSimpleCircuitBreaking() {
+        ByteSizeValue memoryLimitForSimple = enoughMemoryForSimple();
+        Operator.OperatorFactory simple = simple();
+        DriverContext inputFactoryContext = driverContext();
+        List<Page> input = CannedSourceOperator.collectPages(simpleInput(inputFactoryContext.blockFactory(), between(1_000, 10_000)));
+        try {
+            ByteSizeValue limit = BreakerTestUtil.findBreakerLimit(
+                memoryLimitForSimple,
+                l -> runWithLimit(simple, CannedSourceOperator.deepCopyOf(input), l)
+            );
+            ByteSizeValue testWithSize = ByteSizeValue.ofBytes(randomLongBetween(0, limit.getBytes()));
+            logger.info("testing with {} against a limit of {}", testWithSize, limit);
+            Exception e = expectThrows(
+                CircuitBreakingException.class,
+                () -> runWithLimit(simple, CannedSourceOperator.deepCopyOf(input), testWithSize)
+            );
+            assertThat(e.getMessage(), equalTo(MockBigArrays.ERROR_MESSAGE));
+        } finally {
+            Releasables.closeExpectNoException(Releasables.wrap(() -> Iterators.map(input.iterator(), p -> p::releaseBlocks)));
+        }
+        assertThat(inputFactoryContext.breaker().getUsed(), equalTo(0L));
     }
 
-    private void testSimpleCircuitBreaking(ByteSizeValue limit) {
-        /*
-         * We build two CircuitBreakers - one for the input blocks and one for the operation itself.
-         * The input blocks don't count against the memory usage for the limited operator that we
-         * build.
-         */
-        DriverContext inputFactoryContext = driverContext();
+    private void runWithLimit(Operator.OperatorFactory factory, List<Page> input, ByteSizeValue limit) {
         BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, limit).withCircuitBreaking();
-        Operator.OperatorFactory simple = simple();
-        logger.info("running {} with {}", simple, bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST));
-        List<Page> input = CannedSourceOperator.collectPages(simpleInput(inputFactoryContext.blockFactory(), between(1_000, 10_000)));
         CircuitBreaker breaker = bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST);
         BlockFactory blockFactory = BlockFactory.getInstance(breaker, bigArrays);
         DriverContext driverContext = new DriverContext(bigArrays, blockFactory);
-        boolean[] driverStarted = new boolean[1];
-        Exception e = expectThrows(CircuitBreakingException.class, () -> {
-            var operator = simple.get(driverContext);
-            driverStarted[0] = true;
+        boolean driverStarted = false;
+        try {
+            var operator = factory.get(driverContext);
+            driverStarted = true;
             drive(operator, input.iterator(), driverContext);
-        });
-        if (driverStarted[0] == false) {
-            // if drive hasn't even started then we need to release the input pages
-            Releasables.closeExpectNoException(Releasables.wrap(() -> Iterators.map(input.iterator(), p -> p::releaseBlocks)));
+        } finally {
+            if (driverStarted == false) {
+                // if drive hasn't even started then we need to release the input pages manually
+                Releasables.closeExpectNoException(Releasables.wrap(() -> Iterators.map(input.iterator(), p -> p::releaseBlocks)));
+            }
+            assertThat(bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST).getUsed(), equalTo(0L));
         }
-        assertThat(e.getMessage(), equalTo(MockBigArrays.ERROR_MESSAGE));
-        assertThat(bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST).getUsed(), equalTo(0L));
-
-        // Note the lack of try/finally here - we're asserting that when the driver throws an exception we clear the breakers.
-        assertThat(bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST).getUsed(), equalTo(0L));
-        assertThat(inputFactoryContext.breaker().getUsed(), equalTo(0L));
     }
 
     /**

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ProjectOperatorTests.java

@@ -96,7 +96,7 @@ public class ProjectOperatorTests extends OperatorTestCase {
     }
 
     @Override
-    protected ByteSizeValue memoryLimitForSimple() {
+    protected ByteSizeValue enoughMemoryForSimple() {
         assumeTrue("doesn't allocate", false);
         return null;
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.compute.operator;
 
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
@@ -83,11 +82,6 @@ public class StringExtractOperatorTests extends OperatorTestCase {
         }
     }
 
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        return ByteSizeValue.ofKb(15);
-    }
-
     public void testMultivalueDissectInput() {
 
         StringExtractOperator operator = new StringExtractOperator(new String[] { "test" }, new EvalOperator.ExpressionEvaluator() {

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java

@@ -177,15 +177,6 @@ public class TopNOperatorTests extends OperatorTestCase {
         );
     }
 
-    @Override
-    protected ByteSizeValue memoryLimitForSimple() {
-        /*
-         * 775 causes us to blow up while collecting values and 780 doesn't
-         * trip the breaker.
-         */
-        return ByteSizeValue.ofBytes(775);
-    }
-
     public void testRamBytesUsed() {
         RamUsageTester.Accumulator acc = new RamUsageTester.Accumulator() {
             @Override