Double close on shuffle during field loading (#130838)

Fixes a bug during field loading where we could double-close blocks if
we failed to allocate memory during the un-shuffling portion of field
loading from single segments.

Unit test incoming in the followup.

Closes #130426
Closes #130790
Closes #130791
Closes #130792
Closes #130793
Closes #130270
Closes #130788
Closes #130122
Closes #130827

Author: nik9000

muted-tests.yml

@@ -546,39 +546,12 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/130067
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testForPushDownEnrichRule
-  issue: https://github.com/elastic/elasticsearch/issues/130426
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testAvgDurationByArtist
-  issue: https://github.com/elastic/elasticsearch/issues/130790
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testRejection
-  issue: https://github.com/elastic/elasticsearch/issues/130791
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testManyDocuments
-  issue: https://github.com/elastic/elasticsearch/issues/130792
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testListeningRatio
-  issue: https://github.com/elastic/elasticsearch/issues/130793
 - class: org.elasticsearch.xpack.esql.action.EsqlRemoteErrorWrapIT
   method: testThatRemoteErrorsAreWrapped
   issue: https://github.com/elastic/elasticsearch/issues/130794
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testProfile
-  issue: https://github.com/elastic/elasticsearch/issues/130270
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testSumDurationByArtist
-  issue: https://github.com/elastic/elasticsearch/issues/130788
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testTopN
-  issue: https://github.com/elastic/elasticsearch/issues/130122
 - class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT
   method: test {ints.InCast SYNC}
   issue: https://github.com/elastic/elasticsearch/issues/130796
-- class: org.elasticsearch.xpack.esql.action.EnrichIT
-  method: testFilterAfterEnrich
-  issue: https://github.com/elastic/elasticsearch/issues/130827
 - class: org.elasticsearch.backwards.MixedClusterClientYamlTestSuiteIT
   method: test {p0=mtermvectors/10_basic/Tests catching other exceptions per item}
   issue: https://github.com/elastic/elasticsearch/issues/122414

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesFromSingleReader.java

@@ -65,22 +65,27 @@ public int get(int i) {
             return;
         }
         int[] forwards = docs.shardSegmentDocMapForwards();
-        loadFromSingleLeaf(target, new BlockLoader.Docs() {
-            @Override
-            public int count() {
-                return docs.getPositionCount();
-            }
+        Block[] unshuffled = new Block[target.length];
+        try {
+            loadFromSingleLeaf(unshuffled, new BlockLoader.Docs() {
+                @Override
+                public int count() {
+                    return docs.getPositionCount();
+                }
 
-            @Override
-            public int get(int i) {
-                return docs.docs().getInt(forwards[i]);
-            }
-        });
-        final int[] backwards = docs.shardSegmentDocMapBackwards();
-        for (int i = 0; i < target.length; i++) {
-            try (Block in = target[i]) {
-                target[i] = in.filter(backwards);
+                @Override
+                public int get(int i) {
+                    return docs.docs().getInt(forwards[i]);
+                }
+            });
+            final int[] backwards = docs.shardSegmentDocMapBackwards();
+            for (int i = 0; i < unshuffled.length; i++) {
+                target[i] = unshuffled[i].filter(backwards);
+                unshuffled[i].close();
+                unshuffled[i] = null;
             }
+        } finally {
+            Releasables.closeExpectNoException(unshuffled);
         }
     }
 

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/ValuesReader.java

@@ -36,6 +36,9 @@ public Block[] next() {
         boolean success = false;
         try {
             load(target, offset);
+            if (target[0].getPositionCount() != docs.getPositionCount()) {
+                throw new IllegalStateException("partial pages not yet supported");
+            }
             success = true;
             for (Block b : target) {
                 operator.valuesLoaded += b.getTotalValueCount();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/read/ValuesSourceReaderOperatorTests.java

@@ -29,7 +29,6 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
@@ -446,12 +445,6 @@ protected void assertSimpleOutput(List<Page> input, List<Page> results) {
         assertThat(sum, equalTo(expectedSum));
     }
 
-    @Override
-    protected ByteSizeValue enoughMemoryForSimple() {
-        assumeFalse("strange exception in the test, fix soon", true);
-        return ByteSizeValue.ofKb(1);
-    }
-
     public void testLoadAll() {
         DriverContext driverContext = driverContext();
         loadSimpleAndAssert(

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/OperatorTestCase.java

@@ -98,10 +98,16 @@ protected ByteSizeValue enoughMemoryForSimple() {
      * all pages.
      */
     public final void testSimpleCircuitBreaking() {
-        ByteSizeValue memoryLimitForSimple = enoughMemoryForSimple();
-        Operator.OperatorFactory simple = simple(new SimpleOptions(true));
+        /*
+         * Build the input before building `simple` to handle the rare
+         * cases where `simple` need some state from the input - mostly
+         * this is ValuesSourceReaderOperator.
+         */
         DriverContext inputFactoryContext = driverContext();
         List<Page> input = CannedSourceOperator.collectPages(simpleInput(inputFactoryContext.blockFactory(), between(1_000, 10_000)));
+
+        ByteSizeValue memoryLimitForSimple = enoughMemoryForSimple();
+        Operator.OperatorFactory simple = simple(new SimpleOptions(true));
         try {
             ByteSizeValue limit = BreakerTestUtil.findBreakerLimit(memoryLimitForSimple, l -> runWithLimit(simple, input, l));
             ByteSizeValue testWithSize = ByteSizeValue.ofBytes(randomLongBetween(0, limit.getBytes()));

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EnrichIT.java

@@ -50,7 +50,6 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.enrich.EnrichLookupService;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
-import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.junit.After;
 import org.junit.Before;
 
@@ -81,8 +80,8 @@ public class EnrichIT extends AbstractEsqlIntegTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
-        List<Class<? extends Plugin>> plugins = new ArrayList<>(super.nodePlugins());
-        plugins.add(EsqlPlugin.class);
+        List<Class<? extends Plugin>> plugins = new ArrayList<>();
+        plugins.add(EsqlActionBreakerIT.EsqlTestPluginWithMockBlockFactory.class);
         plugins.add(InternalExchangePlugin.class);
         plugins.add(LocalStateEnrich.class);
         plugins.add(IngestCommonPlugin.class);