More test fixes

Author: GalLalouche

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

@@ -82,10 +82,15 @@ public DocVector(
                 "invalid position count [" + shards.getPositionCount() + " != " + docs.getPositionCount() + "]"
             );
         }
-
-        this.uniqueShards = computeUniqueShards(shards);
-        forEachShardRefCounter(RefCounted::mustIncRef);
-        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) {
@@ -359,20 +364,21 @@ public long ramBytesUsed() {
     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
         );
-        forEachShardRefCounter(RefCounted::decRef);
     }
 
     private void forEachShardRefCounter(Consumer<RefCounted> consumer) {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/Driver.java

@@ -317,7 +317,7 @@ private int closeEarlyFinishedOperators(ListIterator<Operator> operators) {
             if (iterator.previous().isFinished()) {
                 var index = iterator.nextIndex();
                 /*
-                 * Remove this operator and all source operators in the
+                 * Close and remove this operator and all source operators in the
                  * most paranoid possible way. Closing operators shouldn't throw,
                  * but if it does, this will make sure we don't try to close any
                  * that succeed twice.

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java

@@ -357,8 +357,8 @@ static final class OrdinalSegmentAggregator implements Releasable, SeenGroupIds
             RefCounted shardRefCounted
         ) throws IOException {
             boolean success = false;
-            shardRefCounted.mustIncRef();
             this.shardRefCounted = shardRefCounted;
+            this.shardRefCounted.mustIncRef();
             List<GroupingAggregator> groupingAggregators = null;
             BitArray bitArray = null;
             try {
@@ -370,7 +370,6 @@ static final class OrdinalSegmentAggregator implements Releasable, SeenGroupIds
                 this.docValuesSupplier = docValuesSupplier;
                 this.aggregators = groupingAggregators;
                 this.visitedOrds = bitArray;
-                this.shardRefCounted.mustIncRef();
                 success = true;
             } finally {
                 if (success == false) {