ESQL: Fix unreleased block in topn (#101648) (#101694)

This fixes a bug in the topn operator where it'll fail to de-track
blocks if there's a failure while building the results from topn.

Author: nik9000

docs/changelog/101648.yaml

@@ -0,0 +1,6 @@
+pr: 101648
+summary: "ESQL: Fix unreleased block in topn"
+area: ES|QL
+type: bug
+issues:
+ - 101588

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
 
 import java.util.List;
 
@@ -207,6 +208,24 @@ interface Builder extends Releasable {
          * Builds the block. This method can be called multiple times.
          */
         Block build();
+
+        /**
+         * Build many {@link Block}s at once, releasing any partially built blocks
+         * if any fail.
+         */
+        static Block[] buildAll(Block.Builder... builders) {
+            Block[] blocks = new Block[builders.length];
+            try {
+                for (int b = 0; b < blocks.length; b++) {
+                    blocks[b] = builders[b].build();
+                }
+            } finally {
+                if (blocks[blocks.length - 1] == null) {
+                    Releasables.closeExpectNoException(blocks);
+                }
+            }
+            return blocks;
+        }
     }
 
     /**

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/TopNOperator.java

@@ -91,12 +91,6 @@ public long ramBytesUsed() {
             return SHALLOW_SIZE + keys.ramBytesUsed() + orderByCompositeKeyAscending.size() / Byte.SIZE + values.ramBytesUsed();
         }
 
-        private void clear() {
-            keys.clear();
-            orderByCompositeKeyAscending.clear();
-            values.clear();
-        }
-
         @Override
         public void close() {
             Releasables.closeExpectNoException(keys, values);
@@ -405,7 +399,17 @@ private Iterator<Page> toPages() {
 
                 p++;
                 if (p == size) {
-                    result.add(new Page(Arrays.stream(builders).map(ResultBuilder::build).toArray(Block[]::new)));
+                    Block[] blocks = new Block[builders.length];
+                    try {
+                        for (int b = 0; b < blocks.length; b++) {
+                            blocks[b] = builders[b].build();
+                        }
+                    } finally {
+                        if (blocks[blocks.length - 1] == null) {
+                            Releasables.closeExpectNoException(blocks);
+                        }
+                    }
+                    result.add(new Page(blocks));
                     Releasables.closeExpectNoException(builders);
                     builders = null;
                 }

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.MockBlockFactory;
+import org.elasticsearch.indices.CrankyCircuitBreakerService;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.test.ESTestCase;
 import org.junit.After;
@@ -107,6 +108,16 @@ protected DriverContext driverContext() { // TODO make this final and return a b
     private final List<CircuitBreaker> breakers = new ArrayList<>();
     private final List<BlockFactory> blockFactories = new ArrayList<>();
 
+    protected final DriverContext crankyDriverContext() {
+        CrankyCircuitBreakerService cranky = new CrankyCircuitBreakerService();
+        BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, cranky).withCircuitBreaking();
+        CircuitBreaker breaker = bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST);
+        breakers.add(breaker);
+        BlockFactory blockFactory = new MockBlockFactory(breaker, bigArrays);
+        blockFactories.add(blockFactory);
+        return new DriverContext(bigArrays, blockFactory);
+    }
+
     /**
      * A {@link DriverContext} with a breaking {@link BigArrays} and {@link BlockFactory}.
      */

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.operator;
 
+import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Tuple;
@@ -44,23 +45,23 @@ public TupleBlockSourceOperator(BlockFactory blockFactory, List<Tuple<Long, Long
 
     @Override
     protected Page createPage(int positionOffset, int length) {
-        var blockBuilder1 = blockFactory.newLongBlockBuilder(length);
-        var blockBuilder2 = blockFactory.newLongBlockBuilder(length);
-        for (int i = 0; i < length; i++) {
-            Tuple<Long, Long> item = values.get(positionOffset + i);
-            if (item.v1() == null) {
-                blockBuilder1.appendNull();
-            } else {
-                blockBuilder1.appendLong(item.v1());
-            }
-            if (item.v2() == null) {
-                blockBuilder2.appendNull();
-            } else {
-                blockBuilder2.appendLong(item.v2());
+        try (var blockBuilder1 = blockFactory.newLongBlockBuilder(length); var blockBuilder2 = blockFactory.newLongBlockBuilder(length)) {
+            for (int i = 0; i < length; i++) {
+                Tuple<Long, Long> item = values.get(positionOffset + i);
+                if (item.v1() == null) {
+                    blockBuilder1.appendNull();
+                } else {
+                    blockBuilder1.appendLong(item.v1());
+                }
+                if (item.v2() == null) {
+                    blockBuilder2.appendNull();
+                } else {
+                    blockBuilder2.appendLong(item.v2());
+                }
             }
+            currentPosition += length;
+            return new Page(Block.Builder.buildAll(blockBuilder1, blockBuilder2));
         }
-        currentPosition += length;
-        return new Page(blockBuilder1.build(), blockBuilder2.build());
     }
 
     @Override

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

@@ -11,6 +11,7 @@
 import org.apache.lucene.tests.util.RamUsageTester;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
@@ -36,6 +37,7 @@
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.operator.TupleBlockSourceOperator;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.indices.CrankyCircuitBreakerService;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.ListMatcher;
 import org.elasticsearch.xpack.versionfield.Version;
@@ -234,15 +236,29 @@ public long accumulateObject(Object o, long shallowSize, Map<Field, Object> fiel
 
     public void testRandomTopN() {
         for (boolean asc : List.of(true, false)) {
-            int limit = randomIntBetween(1, 20);
-            List<Long> inputValues = randomList(0, 5000, ESTestCase::randomLong);
-            Comparator<Long> comparator = asc ? naturalOrder() : reverseOrder();
-            List<Long> expectedValues = inputValues.stream().sorted(comparator).limit(limit).toList();
-            List<Long> outputValues = topNLong(inputValues, limit, asc, false);
-            assertThat(outputValues, equalTo(expectedValues));
+            testRandomTopN(asc, driverContext());
         }
     }
 
+    public void testRandomTopNCranky() {
+        try {
+            testRandomTopN(randomBoolean(), crankyDriverContext());
+            logger.info("cranky didn't break us");
+        } catch (CircuitBreakingException e) {
+            logger.info("broken", e);
+            assertThat(e.getMessage(), equalTo(CrankyCircuitBreakerService.ERROR_MESSAGE));
+        }
+    }
+
+    private void testRandomTopN(boolean asc, DriverContext context) {
+        int limit = randomIntBetween(1, 20);
+        List<Long> inputValues = randomList(0, 5000, ESTestCase::randomLong);
+        Comparator<Long> comparator = asc ? naturalOrder() : reverseOrder();
+        List<Long> expectedValues = inputValues.stream().sorted(comparator).limit(limit).toList();
+        List<Long> outputValues = topNLong(context, inputValues, limit, asc, false);
+        assertThat(outputValues, equalTo(expectedValues));
+    }
+
     public void testBasicTopN() {
         List<Long> values = Arrays.asList(2L, 1L, 4L, null, 5L, 10L, null, 20L, 4L, 100L);
         assertThat(topNLong(values, 1, true, false), equalTo(Arrays.asList(1L)));
@@ -267,8 +283,15 @@ public void testBasicTopN() {
         assertThat(topNLong(values, 100, false, true), equalTo(Arrays.asList(null, null, 100L, 20L, 10L, 5L, 4L, 4L, 2L, 1L)));
     }
 
-    private List<Long> topNLong(List<Long> inputValues, int limit, boolean ascendingOrder, boolean nullsFirst) {
+    private List<Long> topNLong(
+        DriverContext driverContext,
+        List<Long> inputValues,
+        int limit,
+        boolean ascendingOrder,
+        boolean nullsFirst
+    ) {
         return topNTwoColumns(
+            driverContext,
             inputValues.stream().map(v -> tuple(v, 0L)).toList(),
             limit,
             List.of(LONG, LONG),
@@ -277,6 +300,10 @@ private List<Long> topNLong(List<Long> inputValues, int limit, boolean ascending
         ).stream().map(Tuple::v1).toList();
     }
 
+    private List<Long> topNLong(List<Long> inputValues, int limit, boolean ascendingOrder, boolean nullsFirst) {
+        return topNLong(driverContext(), inputValues, limit, ascendingOrder, nullsFirst);
+    }
+
     public void testCompareInts() {
         testCompare(
             new Page(
@@ -422,6 +449,7 @@ public void testTopNTwoColumns() {
         List<Tuple<Long, Long>> values = Arrays.asList(tuple(1L, 1L), tuple(1L, 2L), tuple(null, null), tuple(null, 1L), tuple(1L, null));
         assertThat(
             topNTwoColumns(
+                driverContext(),
                 values,
                 5,
                 List.of(LONG, LONG),
@@ -432,6 +460,7 @@ public void testTopNTwoColumns() {
         );
         assertThat(
             topNTwoColumns(
+                driverContext(),
                 values,
                 5,
                 List.of(LONG, LONG),
@@ -442,6 +471,7 @@ public void testTopNTwoColumns() {
         );
         assertThat(
             topNTwoColumns(
+                driverContext(),
                 values,
                 5,
                 List.of(LONG, LONG),
@@ -613,13 +643,13 @@ public void testCollectAllValues_RandomMultiValues() {
     }
 
     private List<Tuple<Long, Long>> topNTwoColumns(
+        DriverContext driverContext,
         List<Tuple<Long, Long>> inputValues,
         int limit,
         List<ElementType> elementTypes,
         List<TopNEncoder> encoder,
         List<TopNOperator.SortOrder> sortOrders
     ) {
-        DriverContext driverContext = driverContext();
         List<Tuple<Long, Long>> outputValues = new ArrayList<>();
         try (
             Driver driver = new Driver(

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java

@@ -48,7 +48,10 @@ public void ensureBlocksReleased() {
             CircuitBreakerService breakerService = internalCluster().getInstance(CircuitBreakerService.class, node);
             CircuitBreaker reqBreaker = breakerService.getBreaker(CircuitBreaker.REQUEST);
             try {
-                assertBusy(() -> assertThat("Request breaker not reset to 0 on node: " + node, reqBreaker.getUsed(), equalTo(0L)));
+                assertBusy(() -> {
+                    logger.info("running tasks: {}", client().admin().cluster().prepareListTasks().get());
+                    assertThat("Request breaker not reset to 0 on node: " + node, reqBreaker.getUsed(), equalTo(0L));
+                });
             } catch (Exception e) {
                 assertThat("Request breaker not reset to 0 on node: " + node, reqBreaker.getUsed(), equalTo(0L));
             }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -46,6 +46,8 @@
 import org.elasticsearch.compute.operator.topn.TopNOperator.TopNOperatorFactory;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
@@ -111,6 +113,7 @@
  * drivers that are used to execute the given plan.
  */
 public class LocalExecutionPlanner {
+    private static final Logger logger = LogManager.getLogger(LocalExecutionPlanner.class);
 
     private final String sessionId;
     private final CancellableTask parentTask;
@@ -813,14 +816,15 @@ public List<Driver> createDrivers(String sessionId) {
             try {
                 for (DriverFactory df : driverFactories) {
                     for (int i = 0; i < df.driverParallelism.instanceCount; i++) {
+                        logger.trace("building {} {}", i, df);
                         drivers.add(df.driverSupplier.apply(sessionId));
                     }
                 }
                 success = true;
                 return drivers;
             } finally {
                 if (success == false) {
-                    Releasables.close(() -> Releasables.close(drivers));
+                    Releasables.close(Releasables.wrap(drivers));
                 }
             }
         }