ESQL: Make blocks ref counted

docs/changelog/100408.yaml

@@ -0,0 +1,5 @@
+pr: 100408
+summary: "ESQL: Make blocks ref counted"
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -38,7 +38,7 @@ public interface RefCounted {
     void incRef();
 
     /**
-     * Tries to increment the refCount of this instance. This method will return {@code true} iff the refCount was
+     * Tries to increment the refCount of this instance. This method will return {@code true} iff the refCount was successfully incremented.
      *
      * @see #decRef()
      * @see #incRef()

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

@@ -136,10 +136,7 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-        released = true;
+        super.close();
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 }

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

@@ -17,6 +17,9 @@ public final class BooleanVectorBlock extends AbstractVectorBlock implements Boo
 
     private final BooleanVector vector;
 
+    /**
+     * @param vector considered owned by the current block; must not be used in any other {@code Block}
+     */
     BooleanVectorBlock(BooleanVector vector) {
         super(vector.getPositionCount(), vector.blockFactory());
         this.vector = vector;
@@ -72,15 +75,15 @@ public String toString() {
 
     @Override
     public boolean isReleased() {
-        return released || vector.isReleased();
+        return super.isReleased() || vector.isReleased();
     }
 
     @Override
     public void close() {
-        if (released || vector.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (vector.isReleased()) {
+            throw new IllegalStateException("cannot release block [" + this + "] containing already released vector");
         }
-        released = true;
         Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -140,10 +140,7 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-        released = true;
+        super.close();
         blockFactory.adjustBreaker(-ramBytesUsed() + values.bigArraysRamBytesUsed(), true);
         Releasables.closeExpectNoException(values);
     }

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

@@ -18,6 +18,9 @@ public final class BytesRefVectorBlock extends AbstractVectorBlock implements By
 
     private final BytesRefVector vector;
 
+    /**
+     * @param vector considered owned by the current block; must not be used in any other {@code Block}
+     */
     BytesRefVectorBlock(BytesRefVector vector) {
         super(vector.getPositionCount(), vector.blockFactory());
         this.vector = vector;
@@ -73,15 +76,15 @@ public String toString() {
 
     @Override
     public boolean isReleased() {
-        return released || vector.isReleased();
+        return super.isReleased() || vector.isReleased();
     }
 
     @Override
     public void close() {
-        if (released || vector.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (vector.isReleased()) {
+            throw new IllegalStateException("cannot release block [" + this + "] containing already released vector");
         }
-        released = true;
         Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -136,10 +136,7 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-        released = true;
+        super.close();
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 }

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

@@ -17,6 +17,9 @@ public final class DoubleVectorBlock extends AbstractVectorBlock implements Doub
 
     private final DoubleVector vector;
 
+    /**
+     * @param vector considered owned by the current block; must not be used in any other {@code Block}
+     */
     DoubleVectorBlock(DoubleVector vector) {
         super(vector.getPositionCount(), vector.blockFactory());
         this.vector = vector;
@@ -72,15 +75,15 @@ public String toString() {
 
     @Override
     public boolean isReleased() {
-        return released || vector.isReleased();
+        return super.isReleased() || vector.isReleased();
     }
 
     @Override
     public void close() {
-        if (released || vector.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (vector.isReleased()) {
+            throw new IllegalStateException("cannot release block [" + this + "] containing already released vector");
         }
-        released = true;
         Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -136,10 +136,7 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-        released = true;
+        super.close();
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 }

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

@@ -17,6 +17,9 @@ public final class IntVectorBlock extends AbstractVectorBlock implements IntBloc
 
     private final IntVector vector;
 
+    /**
+     * @param vector considered owned by the current block; must not be used in any other {@code Block}
+     */
     IntVectorBlock(IntVector vector) {
         super(vector.getPositionCount(), vector.blockFactory());
         this.vector = vector;
@@ -72,15 +75,15 @@ public String toString() {
 
     @Override
     public boolean isReleased() {
-        return released || vector.isReleased();
+        return super.isReleased() || vector.isReleased();
     }
 
     @Override
     public void close() {
-        if (released || vector.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (vector.isReleased()) {
+            throw new IllegalStateException("cannot release block [" + this + "] containing already released vector");
         }
-        released = true;
         Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -136,10 +136,7 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-        released = true;
+        super.close();
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 }

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

@@ -17,6 +17,9 @@ public final class LongVectorBlock extends AbstractVectorBlock implements LongBl
 
     private final LongVector vector;
 
+    /**
+     * @param vector considered owned by the current block; must not be used in any other {@code Block}
+     */
     LongVectorBlock(LongVector vector) {
         super(vector.getPositionCount(), vector.blockFactory());
         this.vector = vector;
@@ -72,15 +75,15 @@ public String toString() {
 
     @Override
     public boolean isReleased() {
-        return released || vector.isReleased();
+        return super.isReleased() || vector.isReleased();
     }
 
     @Override
     public void close() {
-        if (released || vector.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (vector.isReleased()) {
+            throw new IllegalStateException("cannot release block [" + this + "] containing already released vector");
         }
-        released = true;
         Releasables.closeExpectNoException(vector);
     }
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/PackedValuesBlockHash.java

@@ -236,7 +236,7 @@ public Block[] getKeys() {
                 }
             } finally {
                 if (keyBlocks[keyBlocks.length - 1] == null) {
-                    Releasables.closeExpectNoException(keyBlocks);
+                    Releasables.closeExpectNoException(Block.releaseByDecRef(keyBlocks));
                 }
             }
             return keyBlocks;

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

@@ -12,7 +12,8 @@
 import java.util.BitSet;
 
 abstract class AbstractBlock implements Block {
-
+    private int references = 1;
+    private boolean released = false;
     private final int positionCount;
 
     @Nullable
@@ -23,8 +24,6 @@ abstract class AbstractBlock implements Block {
 
     protected final BlockFactory blockFactory;
 
-    protected boolean released = false;
-
     /**
      * @param positionCount the number of values in this block
      */
@@ -101,4 +100,54 @@ public BlockFactory blockFactory() {
     public boolean isReleased() {
         return released;
     }
+
+    @Override
+    public final void incRef() {
+        if (references <= 0) {
+            throw new IllegalStateException("can't increase refCount on already released block [" + this + "]");
+        }
+        references++;
+    }
+
+    @Override
+    public final boolean tryIncRef() {
+        if (references <= 0) {
+            return false;
+        }
+        references++;
+        return true;
+    }
+
+    @Override
+    public final boolean decRef() {
+        if (references <= 0) {
+            throw new IllegalStateException("can't release already released block [" + this + "]");
+        }
+
+        references--;
+
+        if (references <= 0) {
+            close();
+            return true;
+        }
+        return false;
+    }
+
+    @Override
+    public final boolean hasReferences() {
+        return references >= 1;
+    }
+
+    @Override
+    public void close() {
+        if (released) {
+            throw new IllegalStateException("can't release already released block [" + this + "]");
+        }
+        if (references > 1) {
+            throw new IllegalStateException("can't close block that is still referenced elsewhere [" + this + "]");
+        }
+        released = true;
+        // In case that there is only 1 reference, allow closing without having to call decRef.
+        references = 0;
+    }
 }

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

@@ -11,12 +11,14 @@
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.BlockLoader;
 
 import java.util.List;
 
+//TODO update javadoc
 /**
  * A Block is a columnar representation of homogenous data. It has a position (row) count, and
  * various data retrieval methods for accessing the underlying data that is stored at a given
@@ -28,7 +30,7 @@
  *
  * <p> Block are immutable and can be passed between threads.
  */
-public interface Block extends Accountable, BlockLoader.Block, NamedWriteable, Releasable {
+public interface Block extends Accountable, BlockLoader.Block, NamedWriteable, RefCounted, Releasable {
 
     /**
      * {@return an efficient dense single-value view of this block}.
@@ -57,7 +59,8 @@ public interface Block extends Accountable, BlockLoader.Block, NamedWriteable, R
     /** The block factory associated with this block. */
     BlockFactory blockFactory();
 
-    /** Tells if this block has been released. A block is released by calling its {@link Block#close()} method. */
+    // TODO mention close as well
+    /** Tells if this block has been released. A block is released by calling its {@link Block#decRef()} method. */
     boolean isReleased();
 
     /**
@@ -222,7 +225,7 @@ static Block[] buildAll(Block.Builder... builders) {
                 }
             } finally {
                 if (blocks[blocks.length - 1] == null) {
-                    Releasables.closeExpectNoException(blocks);
+                    Releasables.closeExpectNoException(releaseByDecRef(blocks));
                 }
             }
             return blocks;
@@ -231,7 +234,7 @@ static Block[] buildAll(Block.Builder... builders) {
 
     /**
      * A reference to a {@link Block}. This is {@link Releasable} and
-     * {@link Ref#close closing} it will {@link Block#close release}
+     * {@link Ref#close closing} it will {@link Block#decRef() release}
      * the underlying {@link Block} if it wasn't borrowed from a {@link Page}.
      *
      * The usual way to use this is:
@@ -266,7 +269,7 @@ public boolean floating() {
         @Override
         public void close() {
             if (floating()) {
-                block.close();
+                block.decRef();
             }
         }
     }
@@ -281,4 +284,20 @@ static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
             ConstantNullBlock.ENTRY
         );
     }
+
+    static Releasable releaseByDecRef(Block b) {
+        return () -> {
+            if (b != null) {
+                b.decRef();
+            }
+        };
+    }
+
+    static Releasable[] releaseByDecRef(Block... blocks) {
+        Releasable[] mappedBlocks = new Releasable[blocks.length];
+        for (int i = 0; i < blocks.length; i++) {
+            mappedBlocks[i] = releaseByDecRef(blocks[i]);
+        }
+        return mappedBlocks;
+    }
 }

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

@@ -98,7 +98,7 @@ public static Block[] fromListRow(BlockFactory blockFactory, List<Object> row, i
             return blocks;
         } finally {
             if (success == false) {
-                Releasables.closeExpectNoException(blocks);
+                Releasables.closeExpectNoException(Block.releaseByDecRef(blocks));
             }
         }
     }