Fix esql enrich memory leak (#110450) (#110493)

This PR was reviewed in #109275

Block and Vector use a non-thread-safe RefCounted. Threads that increase
or decrease the references must have a happen-before relationship.
However, this order is not guaranteed in the enrich lookup for the
reference of selectedPositions. The driver can complete the
MergePositionsOperator, which decreases the reference count of
selectedPositions, while the finally block may also decrease it in a
separate thread. These actions occur without a defined happen-before
relationship.

Closes #108532

Author: dnhatn

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupService.java

@@ -30,7 +30,6 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockStreamInput;
 import org.elasticsearch.compute.data.ElementType;
-import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.LocalCircuitBreaker;
 import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
@@ -42,6 +41,7 @@
 import org.elasticsearch.compute.operator.OutputOperator;
 import org.elasticsearch.core.AbstractRefCounted;
 import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.BlockLoader;
 import org.elasticsearch.index.mapper.MappedFieldType;
@@ -249,29 +249,52 @@ private void doLookup(
         ActionListener<Page> listener
     ) {
         Block inputBlock = inputPage.getBlock(0);
-        final IntBlock selectedPositions;
-        if (inputBlock instanceof OrdinalBytesRefBlock ordinalBytesRefBlock) {
-            inputBlock = ordinalBytesRefBlock.getDictionaryVector().asBlock();
-            selectedPositions = ordinalBytesRefBlock.getOrdinalsBlock();
-            selectedPositions.mustIncRef();
-        } else {
-            selectedPositions = IntVector.range(0, inputBlock.getPositionCount(), blockFactory).asBlock();
+        if (inputBlock.areAllValuesNull()) {
+            listener.onResponse(createNullResponse(inputPage.getPositionCount(), extractFields));
+            return;
         }
-        LocalCircuitBreaker localBreaker = null;
+        final List<Releasable> releasables = new ArrayList<>(6);
+        boolean started = false;
         try {
-            if (inputBlock.areAllValuesNull()) {
-                listener.onResponse(createNullResponse(inputPage.getPositionCount(), extractFields));
-                return;
-            }
-            ShardSearchRequest shardSearchRequest = new ShardSearchRequest(shardId, 0, AliasFilter.EMPTY);
-            SearchContext searchContext = searchService.createSearchContext(shardSearchRequest, SearchService.NO_TIMEOUT);
-            listener = ActionListener.runBefore(listener, searchContext::close);
-            localBreaker = new LocalCircuitBreaker(
+            final ShardSearchRequest shardSearchRequest = new ShardSearchRequest(shardId, 0, AliasFilter.EMPTY);
+            final SearchContext searchContext = searchService.createSearchContext(shardSearchRequest, SearchService.NO_TIMEOUT);
+            releasables.add(searchContext);
+            final LocalCircuitBreaker localBreaker = new LocalCircuitBreaker(
                 blockFactory.breaker(),
                 localBreakerSettings.overReservedBytes(),
                 localBreakerSettings.maxOverReservedBytes()
             );
-            DriverContext driverContext = new DriverContext(bigArrays, blockFactory.newChildFactory(localBreaker));
+            releasables.add(localBreaker);
+            final DriverContext driverContext = new DriverContext(bigArrays, blockFactory.newChildFactory(localBreaker));
+            final ElementType[] mergingTypes = new ElementType[extractFields.size()];
+            for (int i = 0; i < extractFields.size(); i++) {
+                mergingTypes[i] = PlannerUtils.toElementType(extractFields.get(i).dataType());
+            }
+            final int[] mergingChannels = IntStream.range(0, extractFields.size()).map(i -> i + 2).toArray();
+            final MergePositionsOperator mergePositionsOperator;
+            if (inputBlock instanceof OrdinalBytesRefBlock ordinalsBytesRefBlock) {
+                inputBlock = ordinalsBytesRefBlock.getDictionaryVector().asBlock();
+                var selectedPositions = ordinalsBytesRefBlock.getOrdinalsBlock();
+                mergePositionsOperator = new MergePositionsOperator(
+                    1,
+                    mergingChannels,
+                    mergingTypes,
+                    selectedPositions,
+                    driverContext.blockFactory()
+                );
+
+            } else {
+                try (var selectedPositions = IntVector.range(0, inputBlock.getPositionCount(), blockFactory).asBlock()) {
+                    mergePositionsOperator = new MergePositionsOperator(
+                        1,
+                        mergingChannels,
+                        mergingTypes,
+                        selectedPositions,
+                        driverContext.blockFactory()
+                    );
+                }
+            }
+            releasables.add(mergePositionsOperator);
             SearchExecutionContext searchExecutionContext = searchContext.getSearchExecutionContext();
             MappedFieldType fieldType = searchExecutionContext.getFieldType(matchField);
             var queryList = switch (matchType) {
@@ -285,57 +308,13 @@ private void doLookup(
                 queryList,
                 searchExecutionContext.getIndexReader()
             );
-            List<Operator> intermediateOperators = new ArrayList<>(extractFields.size() + 2);
-            final ElementType[] mergingTypes = new ElementType[extractFields.size()];
-            // load the fields
-            List<ValuesSourceReaderOperator.FieldInfo> fields = new ArrayList<>(extractFields.size());
-            for (int i = 0; i < extractFields.size(); i++) {
-                NamedExpression extractField = extractFields.get(i);
-                final ElementType elementType = PlannerUtils.toElementType(extractField.dataType());
-                mergingTypes[i] = elementType;
-                EsPhysicalOperationProviders.ShardContext ctx = new EsPhysicalOperationProviders.DefaultShardContext(
-                    0,
-                    searchContext.getSearchExecutionContext(),
-                    searchContext.request().getAliasFilter()
-                );
-                BlockLoader loader = ctx.blockLoader(
-                    extractField instanceof Alias a ? ((NamedExpression) a.child()).name() : extractField.name(),
-                    EsqlDataTypes.isUnsupported(extractField.dataType()),
-                    MappedFieldType.FieldExtractPreference.NONE
-                );
-                fields.add(
-                    new ValuesSourceReaderOperator.FieldInfo(
-                        extractField.name(),
-                        PlannerUtils.toElementType(extractField.dataType()),
-                        shardIdx -> {
-                            if (shardIdx != 0) {
-                                throw new IllegalStateException("only one shard");
-                            }
-                            return loader;
-                        }
-                    )
-                );
-            }
-            intermediateOperators.add(
-                new ValuesSourceReaderOperator(
-                    driverContext.blockFactory(),
-                    fields,
-                    List.of(
-                        new ValuesSourceReaderOperator.ShardContext(
-                            searchContext.searcher().getIndexReader(),
-                            searchContext::newSourceLoader
-                        )
-                    ),
-                    0
-                )
-            );
-            // merging field-values by position
-            final int[] mergingChannels = IntStream.range(0, extractFields.size()).map(i -> i + 2).toArray();
-            intermediateOperators.add(
-                new MergePositionsOperator(1, mergingChannels, mergingTypes, selectedPositions, driverContext.blockFactory())
-            );
+            releasables.add(queryOperator);
+            var extractFieldsOperator = extractFieldsOperator(searchContext, driverContext, extractFields);
+            releasables.add(extractFieldsOperator);
+
             AtomicReference<Page> result = new AtomicReference<>();
             OutputOperator outputOperator = new OutputOperator(List.of(), Function.identity(), result::set);
+            releasables.add(outputOperator);
             Driver driver = new Driver(
                 "enrich-lookup:" + sessionId,
                 System.currentTimeMillis(),
@@ -351,30 +330,69 @@ private void doLookup(
                     inputPage.getPositionCount()
                 ),
                 queryOperator,
-                intermediateOperators,
+                List.of(extractFieldsOperator, mergePositionsOperator),
                 outputOperator,
                 Driver.DEFAULT_STATUS_INTERVAL,
-                localBreaker
+                Releasables.wrap(searchContext, localBreaker)
             );
             task.addListener(() -> {
                 String reason = Objects.requireNonNullElse(task.getReasonCancelled(), "task was cancelled");
                 driver.cancel(reason);
             });
-
             var threadContext = transportService.getThreadPool().getThreadContext();
-            localBreaker = null;
             Driver.start(threadContext, executor, driver, Driver.DEFAULT_MAX_ITERATIONS, listener.map(ignored -> {
                 Page out = result.get();
                 if (out == null) {
                     out = createNullResponse(inputPage.getPositionCount(), extractFields);
                 }
                 return out;
             }));
+            started = true;
         } catch (Exception e) {
             listener.onFailure(e);
         } finally {
-            Releasables.close(selectedPositions, localBreaker);
+            if (started == false) {
+                Releasables.close(releasables);
+            }
+        }
+    }
+
+    private static Operator extractFieldsOperator(
+        SearchContext searchContext,
+        DriverContext driverContext,
+        List<NamedExpression> extractFields
+    ) {
+        EsPhysicalOperationProviders.ShardContext shardContext = new EsPhysicalOperationProviders.DefaultShardContext(
+            0,
+            searchContext.getSearchExecutionContext(),
+            searchContext.request().getAliasFilter()
+        );
+        List<ValuesSourceReaderOperator.FieldInfo> fields = new ArrayList<>(extractFields.size());
+        for (NamedExpression extractField : extractFields) {
+            BlockLoader loader = shardContext.blockLoader(
+                extractField instanceof Alias a ? ((NamedExpression) a.child()).name() : extractField.name(),
+                EsqlDataTypes.isUnsupported(extractField.dataType()),
+                MappedFieldType.FieldExtractPreference.NONE
+            );
+            fields.add(
+                new ValuesSourceReaderOperator.FieldInfo(
+                    extractField.name(),
+                    PlannerUtils.toElementType(extractField.dataType()),
+                    shardIdx -> {
+                        if (shardIdx != 0) {
+                            throw new IllegalStateException("only one shard");
+                        }
+                        return loader;
+                    }
+                )
+            );
         }
+        return new ValuesSourceReaderOperator(
+            driverContext.blockFactory(),
+            fields,
+            List.of(new ValuesSourceReaderOperator.ShardContext(searchContext.searcher().getIndexReader(), searchContext::newSourceLoader)),
+            0
+        );
     }
 
     private Page createNullResponse(int positionCount, List<NamedExpression> extractFields) {