ESQL: Aggressive release of shard contexts

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesSourceReaderBenchmark.java

@@ -40,6 +40,7 @@
 import org.elasticsearch.compute.data.LongVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.lucene.LuceneSourceOperator;
+import org.elasticsearch.compute.lucene.ShardRefCounted;
 import org.elasticsearch.compute.lucene.ValuesSourceReaderOperator;
 import org.elasticsearch.compute.operator.topn.TopNOperator;
 import org.elasticsearch.core.IOUtils;
@@ -477,6 +478,7 @@ private void setupPages() {
                         pages.add(
                             new Page(
                                 new DocVector(
+                                    ShardRefCounted.ALWAYS_REFERENCED,
                                     blockFactory.newConstantIntBlockWith(0, end - begin).asVector(),
                                     blockFactory.newConstantIntBlockWith(ctx.ord, end - begin).asVector(),
                                     docs.build(),
@@ -512,7 +514,14 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                         if (size >= BLOCK_LENGTH) {
                             pages.add(
                                 new Page(
-                                    new DocVector(blockFactory.newConstantIntVector(0, size), leafs.build(), docs.build(), null).asBlock()
+                                    new DocVector(
+
+                                        ShardRefCounted.ALWAYS_REFERENCED,
+                                        blockFactory.newConstantIntVector(0, size),
+                                        leafs.build(),
+                                        docs.build(),
+                                        null
+                                    ).asBlock()
                                 )
                             );
                             docs = blockFactory.newIntVectorBuilder(BLOCK_LENGTH);
@@ -525,6 +534,8 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                     pages.add(
                         new Page(
                             new DocVector(
+
+                                ShardRefCounted.ALWAYS_REFERENCED,
                                 blockFactory.newConstantIntBlockWith(0, size).asVector(),
                                 leafs.build().asBlock().asVector(),
                                 docs.build(),
@@ -551,6 +562,8 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                         pages.add(
                             new Page(
                                 new DocVector(
+
+                                    ShardRefCounted.ALWAYS_REFERENCED,
                                     blockFactory.newConstantIntVector(0, 1),
                                     blockFactory.newConstantIntVector(next.ord, 1),
                                     blockFactory.newConstantIntVector(next.itr.nextInt(), 1),

docs/changelog/129454.yaml

@@ -0,0 +1,5 @@
+pr: 129454
+summary: Aggressive release of shard contexts
+area: ES|QL
+type: enhancement
+issues: []

libs/core/src/main/java/org/elasticsearch/core/Releasables.java

@@ -202,6 +202,11 @@ public boolean equals(Object obj) {
         }
     }
 
+    /** Creates a {@link Releasable} that calls {@link RefCounted#decRef()} when closed. */
+    public static Releasable fromRefCounted(RefCounted refCounted) {
+        return () -> refCounted.decRef();
+    }
+
     private static class ReleaseOnce extends AtomicReference<Releasable> implements Releasable {
         ReleaseOnce(Releasable releasable) {
             super(releasable);

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -115,6 +115,10 @@ public final void close() {
         closeFuture.onResponse(null);
     }
 
+    public final boolean isClosed() {
+        return closeFuture.isDone();
+    }
+
     /**
      * Should be called before executing the main query and after all other parameters have been set.
      */

test/framework/src/main/java/org/elasticsearch/search/MockSearchService.java

@@ -152,7 +152,12 @@ protected SearchContext createContext(
     @Override
     public SearchContext createSearchContext(ShardSearchRequest request, TimeValue timeout) throws IOException {
         SearchContext searchContext = super.createSearchContext(request, timeout);
-        onPutContext.accept(searchContext.readerContext());
+        try {
+            onCreateSearchContext.accept(searchContext);
+        } catch (Exception e) {
+            searchContext.close();
+            throw e;
+        }
         searchContext.addReleasable(() -> onRemoveContext.accept(searchContext.readerContext()));
         return searchContext;
     }

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

@@ -9,6 +9,8 @@
 
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.compute.lucene.ShardRefCounted;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
 
@@ -17,7 +19,7 @@
 /**
  * Wrapper around {@link DocVector} to make a valid {@link Block}.
  */
-public class DocBlock extends AbstractVectorBlock implements Block {
+public class DocBlock extends AbstractVectorBlock implements Block, RefCounted {
 
     private final DocVector vector;
 
@@ -96,6 +98,12 @@ public static class Builder implements Block.Builder {
         private final IntVector.Builder shards;
         private final IntVector.Builder segments;
         private final IntVector.Builder docs;
+        private ShardRefCounted shardRefCounters = ShardRefCounted.ALWAYS_REFERENCED;
+
+        public Builder setShardRefCounted(ShardRefCounted shardRefCounters) {
+            this.shardRefCounters = shardRefCounters;
+            return this;
+        }
 
         private Builder(BlockFactory blockFactory, int estimatedSize) {
             IntVector.Builder shards = null;
@@ -183,7 +191,7 @@ public DocBlock build() {
                 shards = this.shards.build();
                 segments = this.segments.build();
                 docs = this.docs.build();
-                result = new DocVector(shards, segments, docs, null);
+                result = new DocVector(shardRefCounters, shards, segments, docs, null);
                 return result.asBlock();
             } finally {
                 if (result == null) {

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

@@ -10,10 +10,13 @@
 import org.apache.lucene.util.IntroSorter;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.compute.lucene.ShardRefCounted;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
 
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * {@link Vector} where each entry references a lucene document.
@@ -48,8 +51,21 @@ public final class DocVector extends AbstractVector implements Vector {
      */
     private int[] shardSegmentDocMapBackwards;
 
-    public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean singleSegmentNonDecreasing) {
+    private final ShardRefCounted shardRefCounters;
+
+    public ShardRefCounted shardRefCounted() {
+        return shardRefCounters;
+    }
+
+    public DocVector(
+        ShardRefCounted shardRefCounters,
+        IntVector shards,
+        IntVector segments,
+        IntVector docs,
+        Boolean singleSegmentNonDecreasing
+    ) {
         super(shards.getPositionCount(), shards.blockFactory());
+        this.shardRefCounters = shardRefCounters;
         this.shards = shards;
         this.segments = segments;
         this.docs = docs;
@@ -65,10 +81,19 @@ public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean s
             );
         }
         blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
+
+        forEachShardRefCounter(RefCounted::mustIncRef);
     }
 
-    public DocVector(IntVector shards, IntVector segments, IntVector docs, int[] docMapForwards, int[] docMapBackwards) {
-        this(shards, segments, docs, null);
+    public DocVector(
+        ShardRefCounted shardRefCounters,
+        IntVector shards,
+        IntVector segments,
+        IntVector docs,
+        int[] docMapForwards,
+        int[] docMapBackwards
+    ) {
+        this(shardRefCounters, shards, segments, docs, null);
         this.shardSegmentDocMapForwards = docMapForwards;
         this.shardSegmentDocMapBackwards = docMapBackwards;
     }
@@ -238,7 +263,7 @@ public DocVector filter(int... positions) {
             filteredShards = shards.filter(positions);
             filteredSegments = segments.filter(positions);
             filteredDocs = docs.filter(positions);
-            result = new DocVector(filteredShards, filteredSegments, filteredDocs, null);
+            result = new DocVector(shardRefCounters, filteredShards, filteredSegments, filteredDocs, null);
             return result;
         } finally {
             if (result == null) {
@@ -317,5 +342,20 @@ public void closeInternal() {
             segments,
             docs
         );
+        forEachShardRefCounter(RefCounted::decRef);
+    }
+
+    private void forEachShardRefCounter(Consumer<RefCounted> consumer) {
+        switch (shards) {
+            case ConstantIntVector constantIntVector -> consumer.accept(shardRefCounters.get(constantIntVector.getInt(0)));
+            case ConstantNullVector ignored -> {
+                // Noop
+            }
+            default -> {
+                for (int i = 0; i < shards.getPositionCount(); i++) {
+                    consumer.accept(shardRefCounters.get(shards.getInt(i)));
+                }
+            }
+        }
     }
 }

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasables;
 
 import java.io.IOException;
@@ -40,6 +41,7 @@ public class LuceneCountOperator extends LuceneOperator {
     private final LeafCollector leafCollector;
 
     public static class Factory extends LuceneOperator.Factory {
+        private final List<? extends RefCounted> shardRefCounters;
 
         public Factory(
             List<? extends ShardContext> contexts,
@@ -58,11 +60,12 @@ public Factory(
                 false,
                 ScoreMode.COMPLETE_NO_SCORES
             );
+            this.shardRefCounters = contexts;
         }
 
         @Override
         public SourceOperator get(DriverContext driverContext) {
-            return new LuceneCountOperator(driverContext.blockFactory(), sliceQueue, limit);
+            return new LuceneCountOperator(shardRefCounters, driverContext.blockFactory(), sliceQueue, limit);
         }
 
         @Override
@@ -71,8 +74,13 @@ public String describe() {
         }
     }
 
-    public LuceneCountOperator(BlockFactory blockFactory, LuceneSliceQueue sliceQueue, int limit) {
-        super(blockFactory, PAGE_SIZE, sliceQueue);
+    public LuceneCountOperator(
+        List<? extends RefCounted> shardRefCounters,
+        BlockFactory blockFactory,
+        LuceneSliceQueue sliceQueue,
+        int limit
+    ) {
+        super(shardRefCounters, blockFactory, PAGE_SIZE, sliceQueue);
         this.remainingDocs = limit;
         this.leafCollector = new LeafCollector() {
             @Override

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

@@ -108,6 +108,7 @@ public final long evaluate(long value1, long value2) {
         abstract long bytesToLong(byte[] bytes);
     }
 
+    private final List<? extends ShardContext> contexts;
     private final String fieldName;
     private final NumberType numberType;
 
@@ -130,13 +131,14 @@ public LuceneMaxFactory(
             false,
             ScoreMode.COMPLETE_NO_SCORES
         );
+        this.contexts = contexts;
         this.fieldName = fieldName;
         this.numberType = numberType;
     }
 
     @Override
     public SourceOperator get(DriverContext driverContext) {
-        return new LuceneMinMaxOperator(driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MIN_VALUE);
+        return new LuceneMinMaxOperator(contexts, driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MIN_VALUE);
     }
 
     @Override

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.search.MultiValueMode;
 
 import java.io.IOException;
@@ -108,6 +109,7 @@ public final long evaluate(long value1, long value2) {
         abstract long bytesToLong(byte[] bytes);
     }
 
+    private final List<? extends RefCounted> shardRefCounters;
     private final String fieldName;
     private final NumberType numberType;
 
@@ -130,13 +132,22 @@ public LuceneMinFactory(
             false,
             ScoreMode.COMPLETE_NO_SCORES
         );
+        this.shardRefCounters = contexts;
         this.fieldName = fieldName;
         this.numberType = numberType;
     }
 
     @Override
     public SourceOperator get(DriverContext driverContext) {
-        return new LuceneMinMaxOperator(driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MAX_VALUE);
+        return new LuceneMinMaxOperator(
+            shardRefCounters,
+            driverContext.blockFactory(),
+            sliceQueue,
+            fieldName,
+            numberType,
+            limit,
+            Long.MAX_VALUE
+        );
     }
 
     @Override

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

@@ -20,10 +20,12 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.MultiValueMode;
 
 import java.io.IOException;
+import java.util.List;
 
 /**
  * Operator that finds the min or max value of a field using Lucene searches
@@ -65,14 +67,15 @@ sealed interface NumberType permits LuceneMinFactory.NumberType, LuceneMaxFactor
     private final String fieldName;
 
     LuceneMinMaxOperator(
+        List<? extends RefCounted> shardRefCounters,
         BlockFactory blockFactory,
         LuceneSliceQueue sliceQueue,
         String fieldName,
         NumberType numberType,
         int limit,
         long initialResult
     ) {
-        super(blockFactory, PAGE_SIZE, sliceQueue);
+        super(shardRefCounters, blockFactory, PAGE_SIZE, sliceQueue);
         this.remainingDocs = limit;
         this.numberType = numberType;
         this.fieldName = fieldName;