ESQL: Make blocks ref counted

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

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

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

@@ -72,10 +72,9 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 }

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

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

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

@@ -73,10 +73,9 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 }

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

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

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

@@ -72,10 +72,9 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -132,9 +132,9 @@ private void appendValues(StringBuilder sb) {
 
     @Override
     public void close() {
-        if (block.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(block);
         }
-        Releasables.closeExpectNoException(block);
     }
 }

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

@@ -135,9 +135,9 @@ private void appendValues(StringBuilder sb) {
 
     @Override
     public void close() {
-        if (block.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(block);
         }
-        Releasables.closeExpectNoException(block);
     }
 }

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

@@ -132,9 +132,9 @@ private void appendValues(StringBuilder sb) {
 
     @Override
     public void close() {
-        if (block.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(block);
         }
-        Releasables.closeExpectNoException(block);
     }
 }

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

@@ -132,9 +132,9 @@ private void appendValues(StringBuilder sb) {
 
     @Override
     public void close() {
-        if (block.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(block);
         }
-        Releasables.closeExpectNoException(block);
     }
 }

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

@@ -132,9 +132,9 @@ private void appendValues(StringBuilder sb) {
 
     @Override
     public void close() {
-        if (block.isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(block);
         }
-        Releasables.closeExpectNoException(block);
     }
 }

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

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

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

@@ -72,10 +72,9 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 }

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

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

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

@@ -72,10 +72,9 @@ public String toString() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 }

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

@@ -12,7 +12,7 @@
 import java.util.BitSet;
 
 abstract class AbstractBlock implements Block {
-
+    private int references = 1;
     private final int positionCount;
 
     @Nullable
@@ -23,8 +23,6 @@ abstract class AbstractBlock implements Block {
 
     protected final BlockFactory blockFactory;
 
-    protected boolean released = false;
-
     /**
      * @param positionCount the number of values in this block
      */
@@ -99,6 +97,47 @@ public BlockFactory blockFactory() {
 
     @Override
     public boolean isReleased() {
-        return released;
+        return hasReferences() == false;
+    }
+
+    @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() {
+        close();
+
+        return references <= 0;
+    }
+
+    @Override
+    public final boolean hasReferences() {
+        return references >= 1;
+    }
+
+    /**
+     * When overriding this, {@code super.close()} must be called to ensure correct reference counting.
+     * Used in {@link AbstractBlock#decRef()}.
+     */
+    @Override
+    public void close() {
+        if (references <= 0) {
+            throw new IllegalStateException("can't release already released block [" + this + "]");
+        }
+        references--;
     }
 }

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

@@ -10,7 +10,7 @@
 import java.util.Arrays;
 
 abstract class AbstractFilterBlock implements Block {
-
+    private int references = 1;
     protected final int[] positions;
 
     private final Block block;
@@ -125,4 +125,45 @@ protected final boolean assertPosition(int position) {
             : "illegal position, " + position + ", position count:" + getPositionCount();
         return true;
     }
+
+    @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() {
+        close();
+
+        return references <= 0;
+    }
+
+    @Override
+    public final boolean hasReferences() {
+        return references >= 1;
+    }
+
+    /**
+     * When overriding this, {@code super.close()} must be called to ensure correct reference counting.
+     * Used in {@link AbstractFilterBlock#decRef()}.
+     */
+    @Override
+    public void close() {
+        if (references <= 0) {
+            throw new IllegalStateException("can't release already released block [" + this + "]");
+        }
+        references--;
+    }
 }

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

@@ -11,6 +11,7 @@
 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 java.util.List;
@@ -33,7 +34,7 @@
  *
  * <p> Block are immutable and can be passed between threads.
  */
-public interface Block extends Accountable, NamedWriteable, Releasable {
+public interface Block extends Accountable, NamedWriteable, RefCounted, Releasable {
 
     /**
      * {@return an efficient dense single-value view of this block}.

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

@@ -128,11 +128,10 @@ public String toString() {
 
     @Override
     public void close() {
-        if (isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            blockFactory.adjustBreaker(-ramBytesUsed(), true);
         }
-        released = true;
-        blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 
     static class Builder implements Block.Builder {

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

@@ -72,11 +72,10 @@ public long ramBytesUsed() {
 
     @Override
     public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
+        super.close();
+        if (hasReferences() == false) {
+            Releasables.closeExpectNoException(vector);
         }
-        released = true;
-        Releasables.closeExpectNoException(vector);
     }
 
     /**