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,14 @@
 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.BitSet;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * {@link Vector} where each entry references a lucene document.
@@ -29,6 +33,7 @@ public final class DocVector extends AbstractVector implements Vector {
     public static final int SHARD_SEGMENT_DOC_MAP_PER_ROW_OVERHEAD = Integer.BYTES * 2;
 
     private final IntVector shards;
+    private final IntVector uniqueShards;
     private final IntVector segments;
     private final IntVector docs;
 
@@ -48,8 +53,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;
@@ -64,11 +82,50 @@ public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean s
                 "invalid position count [" + shards.getPositionCount() + " != " + docs.getPositionCount() + "]"
             );
         }
-        blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
+        var uniqueShards = computeUniqueShards(shards);
+        try {
+            blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
+            this.uniqueShards = uniqueShards;
+            forEachShardRefCounter(RefCounted::mustIncRef);
+        } catch (Exception e) {
+            Releasables.close(uniqueShards);
+            throw e;
+        }
+    }
+
+    private static IntVector computeUniqueShards(IntVector shards) {
+        switch (shards) {
+            case ConstantIntVector constantIntVector -> {
+                return shards.blockFactory().newConstantIntVector(constantIntVector.getInt(0), 1);
+            }
+            case ConstantNullVector unused -> {
+                return shards.blockFactory().newConstantIntVector(0, 0);
+            }
+            default -> {
+                var seen = new BitSet(128);
+                try (IntVector.Builder uniqueShardsBuilder = shards.blockFactory().newIntVectorBuilder(shards.getPositionCount())) {
+                    for (int p = 0; p < shards.getPositionCount(); p++) {
+                        int shardId = shards.getInt(p);
+                        if (seen.get(shardId) == false) {
+                            seen.set(shardId);
+                            uniqueShardsBuilder.appendInt(shardId);
+                        }
+                    }
+                    return uniqueShardsBuilder.build();
+                }
+            }
+        }
     }
 
-    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 +295,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) {
@@ -287,35 +344,54 @@ private static long ramBytesOrZero(int[] array) {
 
     public static long ramBytesEstimated(
         IntVector shards,
+        IntVector uniqueShards,
         IntVector segments,
         IntVector docs,
         int[] shardSegmentDocMapForwards,
         int[] shardSegmentDocMapBackwards
     ) {
-        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(shards) + RamUsageEstimator.sizeOf(segments) + RamUsageEstimator.sizeOf(docs)
-            + ramBytesOrZero(shardSegmentDocMapForwards) + ramBytesOrZero(shardSegmentDocMapBackwards);
+        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(shards) + RamUsageEstimator.sizeOf(uniqueShards) + RamUsageEstimator.sizeOf(
+            segments
+        ) + RamUsageEstimator.sizeOf(docs) + ramBytesOrZero(shardSegmentDocMapForwards) + ramBytesOrZero(shardSegmentDocMapBackwards);
     }
 
     @Override
     public long ramBytesUsed() {
-        return ramBytesEstimated(shards, segments, docs, shardSegmentDocMapForwards, shardSegmentDocMapBackwards);
+        return ramBytesEstimated(shards, uniqueShards, segments, docs, shardSegmentDocMapForwards, shardSegmentDocMapBackwards);
     }
 
     @Override
     public void allowPassingToDifferentDriver() {
         super.allowPassingToDifferentDriver();
         shards.allowPassingToDifferentDriver();
+        uniqueShards.allowPassingToDifferentDriver();
         segments.allowPassingToDifferentDriver();
         docs.allowPassingToDifferentDriver();
     }
 
     @Override
     public void closeInternal() {
+        forEachShardRefCounter(RefCounted::decRef);
         Releasables.closeExpectNoException(
             () -> blockFactory().adjustBreaker(-BASE_RAM_BYTES_USED - (shardSegmentDocMapForwards == null ? 0 : sizeOfSegmentDocMap())),
             shards,
+            uniqueShards,
             segments,
             docs
         );
     }
+
+    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 < uniqueShards.getPositionCount(); i++) {
+                    consumer.accept(shardRefCounters.get(uniqueShards.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;