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,10 @@ public boolean equals(Object obj) {
         }
     }
 
+    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,6 +10,7 @@
 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.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
 
@@ -48,8 +49,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 shardRefCounters() {
+        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 +79,19 @@ public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean s
             );
         }
         blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
+
+        forEachShardRefCounter(DecOrInc.INC);
     }
 
-    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 +261,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 +340,32 @@ public void closeInternal() {
             segments,
             docs
         );
+        forEachShardRefCounter(DecOrInc.DEC);
+    }
+
+    private enum DecOrInc {
+        DEC,
+        INC;
+
+        void apply(ShardRefCounted counters, int shardId) {
+            switch (this) {
+                case DEC -> counters.get(shardId).decRef();
+                case INC -> counters.get(shardId).mustIncRef();
+            }
+        }
+    }
+
+    private void forEachShardRefCounter(DecOrInc mode) {
+        switch (shards) {
+            case ConstantIntVector constantIntVector -> mode.apply(shardRefCounters, constantIntVector.getInt(0));
+            case ConstantNullVector ignored -> {
+                // Noop
+            }
+            default -> {
+                for (int i = 0; i < shards.getPositionCount(); i++) {
+                    mode.apply(shardRefCounters, 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;
@@ -34,12 +35,14 @@ public class LuceneCountOperator extends LuceneOperator {
 
     private static final int PAGE_SIZE = 1;
 
+    private final List<? extends RefCounted> shardRefCounters;
     private int totalHits = 0;
     private int remainingDocs;
 
     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 +61,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 +75,15 @@ public String describe() {
         }
     }
 
-    public LuceneCountOperator(BlockFactory blockFactory, LuceneSliceQueue sliceQueue, int limit) {
+    public LuceneCountOperator(
+        List<? extends RefCounted> shardRefCounters,
+        BlockFactory blockFactory,
+        LuceneSliceQueue sliceQueue,
+        int limit
+    ) {
         super(blockFactory, PAGE_SIZE, sliceQueue);
+        this.shardRefCounters = shardRefCounters;
+        shardRefCounters.forEach(RefCounted::mustIncRef);
         this.remainingDocs = limit;
         this.leafCollector = new LeafCollector() {
             @Override
@@ -171,4 +182,9 @@ protected Page getCheckedOutput() throws IOException {
     protected void describe(StringBuilder sb) {
         sb.append(", remainingDocs=").append(remainingDocs);
     }
+
+    @Override
+    public void close() {
+        shardRefCounters.forEach(RefCounted::decRef);
+    }
 }

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

@@ -138,7 +138,11 @@ public final Page getOutput() {
     protected abstract Page getCheckedOutput() throws IOException;
 
     @Override
-    public void close() {}
+    public void close() {
+        if (currentScorer != null) {
+            currentScorer.shardContext().decRef();
+        }
+    }
 
     LuceneScorer getCurrentOrLoadNextScorer() {
         while (currentScorer == null || currentScorer.isDone()) {
@@ -161,7 +165,11 @@ LuceneScorer getCurrentOrLoadNextScorer() {
             ) {
                 final Weight weight = currentSlice.weight();
                 processedQueries.add(weight.getQuery());
+                var previousScorer = currentScorer;
                 currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.tags(), leaf);
+                if (previousScorer != null) {
+                    previousScorer.shardContext().decRef();
+                }
             }
             assert currentScorer.maxPosition <= partialLeaf.maxDoc() : currentScorer.maxPosition + ">" + partialLeaf.maxDoc();
             currentScorer.maxPosition = partialLeaf.maxDoc();
@@ -188,6 +196,7 @@ static final class LuceneScorer {
         private Thread executingThread;
 
         LuceneScorer(ShardContext shardContext, Weight weight, List<Object> tags, LeafReaderContext leafReaderContext) {
+            shardContext.incRef();
             this.shardContext = shardContext;
             this.weight = weight;
             this.tags = tags;