More

Author: nik9000

test/external-modules/esql-heap-attack/src/javaRestTest/java/org/elasticsearch/xpack/esql/heap_attack/HeapAttackIT.java

@@ -212,6 +212,11 @@ public void testSortByManyLongsGiantTopNTooMuchMemory() throws IOException {
         assertCircuitBreaks(attempt -> sortBySomeLongsLimit(attempt * 500000));
     }
 
+    public void testStupidTopN() throws IOException {
+        initManyLongs(1); // Doesn't actually matter how much data there is.
+        assertCircuitBreaks(attempt -> sortBySomeLongsLimit(2147483630));
+    }
+
     private static final int MAX_ATTEMPTS = 5;
 
     interface TryCircuitBreaking {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperator.java

@@ -24,19 +24,16 @@
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.breaker.CircuitBreaker;
-import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DocVector;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.DoubleVector;
-import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.SourceOperator;
-import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.sort.SortBuilder;
@@ -168,7 +165,7 @@ public LuceneTopNSourceOperator(
         this.estimatedPerRowSortSize = estimatedPerRowSortSize;
         this.limit = limit;
         this.needsScore = needsScore;
-        breaker.addEstimateBytesAndMaybeBreak(reserveSize(), "esql topn");
+        breaker.addEstimateBytesAndMaybeBreak(reserveSize(), "esql lucene topn");
     }
 
     @Override

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

@@ -294,14 +294,14 @@ public String describe() {
 
     private final BlockFactory blockFactory;
     private final CircuitBreaker breaker;
-    private final Queue inputQueue;
 
     private final int maxPageSize;
 
     private final List<ElementType> elementTypes;
     private final List<TopNEncoder> encoders;
     private final List<SortOrder> sortOrders;
 
+    private Queue inputQueue;
     private Row spare;
     private int spareValuesPreAllocSize = 0;
     private int spareKeysPreAllocSize = 0;
@@ -346,7 +346,8 @@ public TopNOperator(
         this.elementTypes = elementTypes;
         this.encoders = encoders;
         this.sortOrders = sortOrders;
-        this.inputQueue = new Queue(topCount);
+        breaker.addEstimateBytesAndMaybeBreak(Queue.sizeOf(topCount), "esql engine topn");
+        this.inputQueue = new Queue(breaker, topCount);
     }
 
     static int compareRows(Row r1, Row r2) {
@@ -457,6 +458,8 @@ private Iterator<Page> toPages() {
                 list.add(inputQueue.pop());
             }
             Collections.reverse(list);
+            inputQueue.close();
+            inputQueue = null;
 
             int p = 0;
             int size = 0;
@@ -563,19 +566,27 @@ public Page getOutput() {
 
     @Override
     public void close() {
-        /*
-         * If we close before calling finish then spare and inputQueue will be live rows
-         * that need closing. If we close after calling finish then the output iterator
-         * will contain pages of results that have yet to be returned.
-         */
         Releasables.closeExpectNoException(
+            /*
+             * The spare is used during most collections. It's cleared when this Operator
+             * is finish()ed. So it could be null here.
+             */
             spare,
-            inputQueue == null ? null : Releasables.wrap(inputQueue),
+            /*
+             * The inputQueue is a min heap of all live rows. Closing it will close all
+             * the rows it contains and all decrement the breaker for the size of
+             * the heap itself.
+             */
+            inputQueue,
+            /*
+             * If we're in the process of outputting pages then output will contain all
+             * allocated but un-emitted pages.
+             */
             output == null ? null : Releasables.wrap(() -> Iterators.map(output, p -> p::releaseBlocks))
         );
     }
 
-    private static long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TopNOperator.class) + RamUsageEstimator
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TopNOperator.class) + RamUsageEstimator
         .shallowSizeOfInstance(List.class) * 3;
 
     @Override
@@ -589,7 +600,9 @@ public long ramBytesUsed() {
         size += RamUsageEstimator.alignObjectSize(arrHeader + ref * encoders.size());
         size += RamUsageEstimator.alignObjectSize(arrHeader + ref * sortOrders.size());
         size += sortOrders.size() * SortOrder.SHALLOW_SIZE;
-        size += inputQueue.ramBytesUsed();
+        if (inputQueue != null) {
+            size += inputQueue.ramBytesUsed();
+        }
         return size;
     }
 
@@ -598,7 +611,7 @@ public Status status() {
         return new TopNOperatorStatus(
             receiveNanos,
             emitNanos,
-            inputQueue.size(),
+            inputQueue != null ? inputQueue.size() : 0,
             ramBytesUsed(),
             pagesReceived,
             pagesEmitted,
@@ -620,16 +633,14 @@ public String toString() {
             + "]";
     }
 
-    CircuitBreaker breaker() {
-        return breaker;
-    }
-
-    private static class Queue extends PriorityQueue<Row> implements Accountable {
+    private static class Queue extends PriorityQueue<Row> implements Accountable, Releasable {
         private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Queue.class);
+        private final CircuitBreaker breaker;
         private final int maxSize;
 
-        Queue(int maxSize) {
+        Queue(CircuitBreaker breaker, int maxSize) {
             super(maxSize);
+            this.breaker = breaker;
             this.maxSize = maxSize;
         }
 
@@ -654,5 +665,19 @@ public long ramBytesUsed() {
             }
             return total;
         }
+
+        @Override
+        public void close() {
+            Releasables.close(Releasables.wrap(this), () -> breaker.addWithoutBreaking(-Queue.sizeOf(maxSize)));
+
+        }
+
+        public static long sizeOf(int topCount) {
+            long total = SHALLOW_SIZE;
+            total += RamUsageEstimator.alignObjectSize(
+                RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + RamUsageEstimator.NUM_BYTES_OBJECT_REF * ((long) topCount + 1)
+            );
+            return total;
+        }
     }
 }

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

@@ -1489,7 +1489,9 @@ public void testRowResizes() {
             block.decRef();
             op.addInput(new Page(blocks));
 
-            assertThat(breaker.getMemoryRequestCount(), is(94L));
+            // 94 are from the collection process
+            // 1 is for the min-heap itself
+            assertThat(breaker.getMemoryRequestCount(), is(95L));
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java

@@ -69,6 +69,12 @@ public interface Sort {
         Order.OrderDirection direction();
 
         FieldAttribute field();
+
+        /**
+         * Type of the <strong>result</strong> of the sort. For example,
+         * geo distance will be {@link DataType#DOUBLE}.
+         */
+        DataType resulType();
     }
 
     public record FieldSort(FieldAttribute field, Order.OrderDirection direction, Order.NullsPosition nulls) implements Sort {
@@ -80,6 +86,11 @@ public SortBuilder<?> sortBuilder() {
             builder.unmappedType(field.dataType().esType());
             return builder;
         }
+
+        @Override
+        public DataType resulType() {
+            return field.dataType();
+        }
     }
 
     public record GeoDistanceSort(FieldAttribute field, Order.OrderDirection direction, double lat, double lon) implements Sort {
@@ -89,6 +100,11 @@ public SortBuilder<?> sortBuilder() {
             builder.order(Direction.from(direction).asOrder());
             return builder;
         }
+
+        @Override
+        public DataType resulType() {
+            return DataType.DOUBLE;
+        }
     }
 
     public record ScoreSort(Order.OrderDirection direction) implements Sort {
@@ -102,6 +118,11 @@ public FieldAttribute field() {
             // TODO: refactor this: not all Sorts are backed by FieldAttributes
             return null;
         }
+
+        @Override
+        public DataType resulType() {
+            return DataType.DOUBLE;
+        }
     }
 
     public record QueryBuilderAndTags(QueryBuilder query, List<Object> tags) {

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

@@ -298,7 +298,7 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec,
             long estimatedPerRowSortSize = 0;
             for (Sort sort : sorts) {
                 sortBuilders.add(sort.sortBuilder());
-                estimatedPerRowSortSize += EstimatesRowSize.estimateSize(sort.field().dataType());
+                estimatedPerRowSortSize += EstimatesRowSize.estimateSize(sort.resulType());
             }
             /*
              * In the worst case Lucene's TopN keeps each value in memory twice. Once