ESQL: Make vectors RefCounted (#103645)

* Factor refcounting/releasing logic into own class
* Make vectors RefCounted

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

@@ -64,11 +64,7 @@ public final class BooleanBigArrayVector extends AbstractVector implements Boole
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         values.close();
     }
 

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

@@ -86,11 +86,7 @@ final class BytesRefArrayVector extends AbstractVector implements BytesRefVector
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         blockFactory().adjustBreaker(-ramBytesUsed() + values.bigArraysRamBytesUsed(), true);
         Releasables.closeExpectNoException(values);
     }

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

@@ -70,13 +70,4 @@ final class ConstantBooleanVector extends AbstractVector implements BooleanVecto
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

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

@@ -75,13 +75,4 @@ final class ConstantBytesRefVector extends AbstractVector implements BytesRefVec
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

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

@@ -70,13 +70,4 @@ final class ConstantDoubleVector extends AbstractVector implements DoubleVector
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

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

@@ -70,13 +70,4 @@ final class ConstantIntVector extends AbstractVector implements IntVector {
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

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

@@ -70,13 +70,4 @@ final class ConstantLongVector extends AbstractVector implements LongVector {
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

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

@@ -62,11 +62,7 @@ public final class DoubleBigArrayVector extends AbstractVector implements Double
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         values.close();
     }
 

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

@@ -62,11 +62,7 @@ public final class IntBigArrayVector extends AbstractVector implements IntVector
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         values.close();
     }
 

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

@@ -62,11 +62,7 @@ public final class LongBigArrayVector extends AbstractVector implements LongVect
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         values.close();
     }
 

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

@@ -11,8 +11,7 @@ import org.elasticsearch.core.Nullable;
 
 import java.util.BitSet;
 
-abstract class AbstractBlock implements Block {
-    private int references = 1;
+abstract class AbstractBlock extends AbstractNonThreadSafeRefCounted implements Block {
     private final int positionCount;
 
     @Nullable
@@ -104,52 +103,4 @@ abstract class AbstractBlock implements Block {
     public final boolean isReleased() {
         return hasReferences() == false;
     }
-
-    @Override
-    public final void incRef() {
-        if (isReleased()) {
-            throw new IllegalStateException("can't increase refCount on already released block [" + this + "]");
-        }
-        references++;
-    }
-
-    @Override
-    public final boolean tryIncRef() {
-        if (isReleased()) {
-            return false;
-        }
-        references++;
-        return true;
-    }
-
-    @Override
-    public final boolean decRef() {
-        if (isReleased()) {
-            throw new IllegalStateException("can't release already released block [" + this + "]");
-        }
-
-        references--;
-
-        if (references <= 0) {
-            closeInternal();
-            return true;
-        }
-        return false;
-    }
-
-    @Override
-    public final boolean hasReferences() {
-        return references >= 1;
-    }
-
-    @Override
-    public final void close() {
-        decRef();
-    }
-
-    /**
-     * This is called when the number of references reaches zero.
-     * It must release any resources held by the block (adjusting circuit breakers if needed).
-     */
-    protected abstract void closeInternal();
 }

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.data;
+
+import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
+
+/**
+ * Releasable, non-threadsafe version of {@link org.elasticsearch.core.AbstractRefCounted}.
+ * Calls to {@link AbstractNonThreadSafeRefCounted#decRef()} and {@link AbstractNonThreadSafeRefCounted#close()} are equivalent.
+ */
+abstract class AbstractNonThreadSafeRefCounted implements RefCounted, Releasable {
+    private int references = 1;
+
+    @Override
+    public final void incRef() {
+        if (hasReferences() == false) {
+            throw new IllegalStateException("can't increase refCount on already released object [" + this + "]");
+        }
+        references++;
+    }
+
+    @Override
+    public final boolean tryIncRef() {
+        if (hasReferences() == false) {
+            return false;
+        }
+        references++;
+        return true;
+    }
+
+    @Override
+    public final boolean decRef() {
+        if (hasReferences() == false) {
+            throw new IllegalStateException("can't release already released object [" + this + "]");
+        }
+
+        references--;
+
+        if (references <= 0) {
+            closeInternal();
+            return true;
+        }
+        return false;
+    }
+
+    @Override
+    public final boolean hasReferences() {
+        return references >= 1;
+    }
+
+    @Override
+    public final void close() {
+        decRef();
+    }
+
+    /**
+     * This is called when the number of references reaches zero.
+     * This is where resources should be released (adjusting circuit breakers if needed).
+     */
+    protected abstract void closeInternal();
+}

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

@@ -10,11 +10,10 @@ package org.elasticsearch.compute.data;
 /**
  * A dense Vector of single values.
  */
-abstract class AbstractVector implements Vector {
+abstract class AbstractVector extends AbstractNonThreadSafeRefCounted implements Vector {
 
     private final int positionCount;
     private BlockFactory blockFactory;
-    protected boolean released;
 
     protected AbstractVector(int positionCount, BlockFactory blockFactory) {
         this.positionCount = positionCount;
@@ -41,16 +40,12 @@ abstract class AbstractVector implements Vector {
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    protected void closeInternal() {
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 
     @Override
     public final boolean isReleased() {
-        return released;
+        return hasReferences() == false;
     }
 }

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

@@ -228,8 +228,7 @@ public final class DocVector extends AbstractVector implements Vector {
     }
 
     @Override
-    public void close() {
-        released = true;
+    public void closeInternal() {
         Releasables.closeExpectNoException(shards.asBlock(), segments.asBlock(), docs.asBlock()); // Ugh! we always close blocks
     }
 }

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

@@ -8,12 +8,13 @@
 package org.elasticsearch.compute.data;
 
 import org.apache.lucene.util.Accountable;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 
 /**
  * A dense Vector of single values.
  */
-public interface Vector extends Accountable, Releasable {
+public interface Vector extends Accountable, RefCounted, Releasable {
 
     /**
      * {@return Returns a new Block containing this vector.}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayVector.java.st

@@ -117,11 +117,7 @@ $endif$
 
 $if(BytesRef)$
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         blockFactory().adjustBreaker(-ramBytesUsed() + values.bigArraysRamBytesUsed(), true);
         Releasables.closeExpectNoException(values);
     }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BigArrayVector.java.st

@@ -71,11 +71,7 @@ public final class $Type$BigArrayVector extends AbstractVector implements $Type$
     }
 
     @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
+    public void closeInternal() {
         values.close();
     }
 

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ConstantVector.java.st

@@ -94,13 +94,4 @@ $endif$
     public String toString() {
         return getClass().getSimpleName() + "[positions=" + getPositionCount() + ", value=" + value + ']';
     }
-
-    @Override
-    public void close() {
-        if (released) {
-            throw new IllegalStateException("can't release already released vector [" + this + "]");
-        }
-        released = true;
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
-    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicBlockTests.java

@@ -11,9 +11,15 @@ import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.common.util.DoubleArray;
+import org.elasticsearch.common.util.IntArray;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
+import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.test.ESTestCase;
@@ -1000,7 +1006,7 @@ public class BasicBlockTests extends ESTestCase {
 
     static void assertCannotDoubleRelease(Block block) {
         var ex = expectThrows(IllegalStateException.class, () -> block.close());
-        assertThat(ex.getMessage(), containsString("can't release already released block"));
+        assertThat(ex.getMessage(), containsString("can't release already released object"));
     }
 
     static void assertCannotReadFromPage(Page page) {
@@ -1035,6 +1041,13 @@ public class BasicBlockTests extends ESTestCase {
         assertThat(breaker.getUsed(), is(0L));
     }
 
+    public void testRefCountingBigArrayBlock() {
+        Block block = randomBigArrayBlock();
+        assertThat(breaker.getUsed(), greaterThan(0L));
+        assertRefCountingBehavior(block);
+        assertThat(breaker.getUsed(), is(0L));
+    }
+
     public void testRefCountingConstantNullBlock() {
         Block block = blockFactory.newConstantNullBlock(10);
         assertThat(breaker.getUsed(), greaterThan(0L));
@@ -1051,52 +1064,165 @@ public class BasicBlockTests extends ESTestCase {
     }
 
     public void testRefCountingVectorBlock() {
-        Block block = randomNonDocVector().asBlock();
+        Block block = randomConstantVector().asBlock();
         assertThat(breaker.getUsed(), greaterThan(0L));
         assertRefCountingBehavior(block);
         assertThat(breaker.getUsed(), is(0L));
     }
 
-    // Take a block with exactly 1 reference and assert that ref counting works fine.
-    static void assertRefCountingBehavior(Block b) {
-        assertTrue(b.hasReferences());
+    public void testRefCountingArrayVector() {
+        Vector vector = randomArrayVector();
+        assertThat(breaker.getUsed(), greaterThan(0L));
+        assertRefCountingBehavior(vector);
+        assertThat(breaker.getUsed(), is(0L));
+    }
+
+    public void testRefCountingBigArrayVector() {
+        Vector vector = randomBigArrayVector();
+        assertThat(breaker.getUsed(), greaterThan(0L));
+        assertRefCountingBehavior(vector);
+        assertThat(breaker.getUsed(), is(0L));
+    }
+
+    public void testRefCountingConstantVector() {
+        Vector vector = randomConstantVector();
+        assertThat(breaker.getUsed(), greaterThan(0L));
+        assertRefCountingBehavior(vector);
+        assertThat(breaker.getUsed(), is(0L));
+    }
+
+    public void testRefCountingDocVector() {
+        int positionCount = randomIntBetween(0, 100);
+        DocVector vector = new DocVector(intVector(positionCount), intVector(positionCount), intVector(positionCount), true);
+        assertThat(breaker.getUsed(), greaterThan(0L));
+        assertRefCountingBehavior(vector);
+        assertThat(breaker.getUsed(), is(0L));
+    }
+
+    /**
+     * Take an object with exactly 1 reference and assert that ref counting works fine.
+     * Assumes that {@link Releasable#close()} and {@link RefCounted#decRef()} are equivalent.
+     */
+    static <T extends RefCounted & Releasable> void assertRefCountingBehavior(T object) {
+        assertTrue(object.hasReferences());
         int numShallowCopies = randomIntBetween(0, 15);
         for (int i = 0; i < numShallowCopies; i++) {
             if (randomBoolean()) {
-                b.incRef();
+                object.incRef();
             } else {
-                assertTrue(b.tryIncRef());
+                assertTrue(object.tryIncRef());
             }
         }
 
         for (int i = 0; i < numShallowCopies; i++) {
             if (randomBoolean()) {
-                b.close();
+                object.close();
             } else {
                 // closing and decRef'ing must be equivalent
-                assertFalse(b.decRef());
+                assertFalse(object.decRef());
             }
-            assertTrue(b.hasReferences());
+            assertTrue(object.hasReferences());
         }
 
         if (randomBoolean()) {
-            b.close();
+            object.close();
         } else {
-            assertTrue(b.decRef());
+            assertTrue(object.decRef());
         }
 
-        assertFalse(b.hasReferences());
-        assertFalse(b.tryIncRef());
+        assertFalse(object.hasReferences());
+        assertFalse(object.tryIncRef());
 
-        expectThrows(IllegalStateException.class, b::close);
-        expectThrows(IllegalStateException.class, b::incRef);
+        expectThrows(IllegalStateException.class, object::close);
+        expectThrows(IllegalStateException.class, object::incRef);
     }
 
     private IntVector intVector(int positionCount) {
         return blockFactory.newIntArrayVector(IntStream.range(0, positionCount).toArray(), positionCount);
     }
 
-    private Vector randomNonDocVector() {
+    private Vector randomArrayVector() {
+        int positionCount = randomIntBetween(0, 100);
+        int vectorType = randomIntBetween(0, 4);
+
+        return switch (vectorType) {
+            case 0 -> {
+                boolean[] values = new boolean[positionCount];
+                Arrays.fill(values, randomBoolean());
+                yield blockFactory.newBooleanArrayVector(values, positionCount);
+            }
+            case 1 -> {
+                BytesRefArray values = new BytesRefArray(positionCount, BigArrays.NON_RECYCLING_INSTANCE);
+                for (int i = 0; i < positionCount; i++) {
+                    values.append(new BytesRef(randomByteArrayOfLength(between(1, 20))));
+                }
+
+                yield blockFactory.newBytesRefArrayVector(values, positionCount);
+            }
+            case 2 -> {
+                double[] values = new double[positionCount];
+                Arrays.fill(values, 1.0);
+
+                yield blockFactory.newDoubleArrayVector(values, positionCount);
+            }
+            case 3 -> {
+                int[] values = new int[positionCount];
+                Arrays.fill(values, 1);
+
+                yield blockFactory.newIntArrayVector(values, positionCount);
+            }
+            default -> {
+                long[] values = new long[positionCount];
+                Arrays.fill(values, 1L);
+
+                yield blockFactory.newLongArrayVector(values, positionCount);
+            }
+        };
+    }
+
+    private Vector randomBigArrayVector() {
+        int positionCount = randomIntBetween(0, 10000);
+        int arrayType = randomIntBetween(0, 3);
+
+        return switch (arrayType) {
+            case 0 -> {
+                BitArray values = new BitArray(positionCount, blockFactory.bigArrays());
+                for (int i = 0; i < positionCount; i++) {
+                    if (randomBoolean()) {
+                        values.set(positionCount);
+                    }
+                }
+
+                yield new BooleanBigArrayVector(values, positionCount, blockFactory);
+            }
+            case 1 -> {
+                DoubleArray values = blockFactory.bigArrays().newDoubleArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomDouble());
+                }
+
+                yield new DoubleBigArrayVector(values, positionCount, blockFactory);
+            }
+            case 2 -> {
+                IntArray values = blockFactory.bigArrays().newIntArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomInt());
+                }
+
+                yield new IntBigArrayVector(values, positionCount, blockFactory);
+            }
+            default -> {
+                LongArray values = blockFactory.bigArrays().newLongArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomLong());
+                }
+
+                yield new LongBigArrayVector(values, positionCount, blockFactory);
+            }
+        };
+    }
+
+    private Vector randomConstantVector() {
         int positionCount = randomIntBetween(0, 100);
         int vectorType = randomIntBetween(0, 4);
 
@@ -1116,7 +1242,7 @@ public class BasicBlockTests extends ESTestCase {
         return switch (arrayType) {
             case 0 -> {
                 boolean[] values = new boolean[positionCount];
-                Arrays.fill(values, true);
+                Arrays.fill(values, randomBoolean());
 
                 yield blockFactory.newBooleanArrayBlock(values, positionCount, new int[] {}, new BitSet(), randomOrdering());
             }
@@ -1148,4 +1274,46 @@ public class BasicBlockTests extends ESTestCase {
             }
         };
     }
+
+    private Block randomBigArrayBlock() {
+        int positionCount = randomIntBetween(0, 10000);
+        int arrayType = randomIntBetween(0, 3);
+
+        return switch (arrayType) {
+            case 0 -> {
+                BitArray values = new BitArray(positionCount, blockFactory.bigArrays());
+                for (int i = 0; i < positionCount; i++) {
+                    if (randomBoolean()) {
+                        values.set(positionCount);
+                    }
+                }
+
+                yield new BooleanBigArrayBlock(values, positionCount, null, new BitSet(), Block.MvOrdering.UNORDERED, blockFactory);
+            }
+            case 1 -> {
+                DoubleArray values = blockFactory.bigArrays().newDoubleArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomDouble());
+                }
+
+                yield new DoubleBigArrayBlock(values, positionCount, null, new BitSet(), Block.MvOrdering.UNORDERED, blockFactory);
+            }
+            case 2 -> {
+                IntArray values = blockFactory.bigArrays().newIntArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomInt());
+                }
+
+                yield new IntBigArrayBlock(values, positionCount, null, new BitSet(), Block.MvOrdering.UNORDERED, blockFactory);
+            }
+            default -> {
+                LongArray values = blockFactory.bigArrays().newLongArray(positionCount, false);
+                for (int i = 0; i < positionCount; i++) {
+                    values.set(i, randomLong());
+                }
+
+                yield new LongBigArrayBlock(values, positionCount, null, new BitSet(), Block.MvOrdering.UNORDERED, blockFactory);
+            }
+        };
+    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/DocVectorTests.java

@@ -145,7 +145,7 @@ public class DocVectorTests extends ComputeTestCase {
         assertThat(block.isReleased(), is(true));
 
         Exception e = expectThrows(IllegalStateException.class, () -> block.close());
-        assertThat(e.getMessage(), containsString("can't release already released block"));
+        assertThat(e.getMessage(), containsString("can't release already released object"));
 
         e = expectThrows(IllegalStateException.class, () -> page.getBlock(0));
         assertThat(e.getMessage(), containsString("can't read released block"));