ESQL: Unit tests for LOOKUP (#120559) (#120719)

This adds a unit test for LOOKUP that's going to be quite good at
finding leaks.

Author: nik9000

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperServiceTestCase.java

@@ -196,7 +196,7 @@ protected final MapperService createMapperService(BooleanSupplier idFieldEnabled
         return createMapperService(getVersion(), getIndexSettings(), idFieldEnabled, mappings);
     }
 
-    protected final MapperService createMapperService(String mappings) throws IOException {
+    public final MapperService createMapperService(String mappings) throws IOException {
         MapperService mapperService = createMapperService(mapping(b -> {}));
         merge(mapperService, mappings);
         return mapperService;
@@ -744,7 +744,7 @@ protected SearchExecutionContext createSearchExecutionContext(MapperService mapp
         return createSearchExecutionContext(mapperService, null, Settings.EMPTY);
     }
 
-    protected SearchExecutionContext createSearchExecutionContext(MapperService mapperService, IndexSearcher searcher) {
+    public final SearchExecutionContext createSearchExecutionContext(MapperService mapperService, IndexSearcher searcher) {
         return createSearchExecutionContext(mapperService, searcher, Settings.EMPTY);
     }
 

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

@@ -134,4 +134,9 @@ public void close() {
             breaker.addWithoutBreaking(-reservedBytes);
         }
     }
+
+    @Override
+    public String toString() {
+        return "LocalCircuitBreaker[" + reservedBytes + "/" + overReservedBytes + ":" + maxOverReservedBytes + "]";
+    }
 }

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

@@ -241,7 +241,7 @@ protected final void assertSimple(DriverContext context, int size) {
      * Tests that finish then close without calling {@link Operator#getOutput} to
      * retrieve a potential last page, releases all memory.
      */
-    public void testSimpleFinishClose() {
+    public void testSimpleFinishClose() throws Exception {
         DriverContext driverContext = driverContext();
         List<Page> input = CannedSourceOperator.collectPages(simpleInput(driverContext.blockFactory(), 1));
         assert input.size() == 1 : "Expected single page, got: " + input;
@@ -253,7 +253,6 @@ public void testSimpleFinishClose() {
             operator.addInput(page);
             operator.finish();
         }
-        assertThat(driverContext.blockFactory().breaker().getUsed(), equalTo(0L));
     }
 
     protected final List<Page> drive(Operator operator, Iterator<Page> input, DriverContext driverContext) {

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

@@ -210,12 +210,13 @@ private void runLookup(PopulateIndices populateIndices) throws IOException {
                 new AsyncExecutionId("test", TaskId.EMPTY_TASK_ID),
                 TEST_REQUEST_TIMEOUT
             );
+            final String finalNodeWithShard = nodeWithShard;
             LookupFromIndexOperator.Factory lookup = new LookupFromIndexOperator.Factory(
                 "test",
                 parentTask,
                 QueryPragmas.ENRICH_MAX_WORKERS.get(Settings.EMPTY),
                 1,
-                internalCluster().getInstance(TransportEsqlQueryAction.class, nodeWithShard).getLookupFromIndexService(),
+                ctx -> internalCluster().getInstance(TransportEsqlQueryAction.class, finalNodeWithShard).getLookupFromIndexService(),
                 DataType.KEYWORD,
                 "lookup",
                 "data",

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

@@ -133,7 +133,7 @@
 public abstract class AbstractLookupService<R extends AbstractLookupService.Request, T extends AbstractLookupService.TransportRequest> {
     private final String actionName;
     private final ClusterService clusterService;
-    private final CreateShardContext createShardContext;
+    private final LookupShardContextFactory lookupShardContextFactory;
     private final TransportService transportService;
     private final Executor executor;
     private final BigArrays bigArrays;
@@ -152,7 +152,7 @@ public abstract class AbstractLookupService<R extends AbstractLookupService.Requ
     AbstractLookupService(
         String actionName,
         ClusterService clusterService,
-        CreateShardContext createShardContext,
+        LookupShardContextFactory lookupShardContextFactory,
         TransportService transportService,
         BigArrays bigArrays,
         BlockFactory blockFactory,
@@ -161,7 +161,7 @@ public abstract class AbstractLookupService<R extends AbstractLookupService.Requ
     ) {
         this.actionName = actionName;
         this.clusterService = clusterService;
-        this.createShardContext = createShardContext;
+        this.lookupShardContextFactory = lookupShardContextFactory;
         this.transportService = transportService;
         this.executor = transportService.getThreadPool().executor(ThreadPool.Names.SEARCH);
         this.bigArrays = bigArrays;
@@ -327,7 +327,7 @@ private void doLookup(T request, CancellableTask task, ActionListener<List<Page>
         final List<Releasable> releasables = new ArrayList<>(6);
         boolean started = false;
         try {
-            LookupShardContext shardContext = createShardContext.create(request.shardId);
+            LookupShardContext shardContext = lookupShardContextFactory.create(request.shardId);
             releasables.add(shardContext.release);
             final LocalCircuitBreaker localBreaker = new LocalCircuitBreaker(
                 blockFactory.breaker(),
@@ -680,10 +680,10 @@ public boolean hasReferences() {
     /**
      * Create a {@link LookupShardContext} for a locally allocated {@link ShardId}.
      */
-    public interface CreateShardContext {
+    public interface LookupShardContextFactory {
         LookupShardContext create(ShardId shardId) throws IOException;
 
-        static CreateShardContext fromSearchService(SearchService searchService) {
+        static LookupShardContextFactory fromSearchService(SearchService searchService) {
             return shardId -> {
                 ShardSearchRequest shardSearchRequest = new ShardSearchRequest(shardId, 0, AliasFilter.EMPTY);
                 return LookupShardContext.fromSearchContext(

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

@@ -47,15 +47,15 @@ public class EnrichLookupService extends AbstractLookupService<EnrichLookupServi
 
     public EnrichLookupService(
         ClusterService clusterService,
-        CreateShardContext createShardContext,
+        LookupShardContextFactory lookupShardContextFactory,
         TransportService transportService,
         BigArrays bigArrays,
         BlockFactory blockFactory
     ) {
         super(
             LOOKUP_ACTION_NAME,
             clusterService,
-            createShardContext,
+            lookupShardContextFactory,
             transportService,
             bigArrays,
             blockFactory,

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

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Function;
 
 // TODO rename package
 public final class LookupFromIndexOperator extends AsyncOperator<LookupFromIndexOperator.OngoingJoin> {
@@ -40,7 +41,7 @@ public record Factory(
         CancellableTask parentTask,
         int maxOutstandingRequests,
         int inputChannel,
-        LookupFromIndexService lookupService,
+        Function<DriverContext, LookupFromIndexService> lookupService,
         DataType inputDataType,
         String lookupIndex,
         String matchField,
@@ -51,6 +52,8 @@ public record Factory(
         public String describe() {
             return "LookupOperator[index="
                 + lookupIndex
+                + " input_type="
+                + inputDataType
                 + " match_field="
                 + matchField
                 + " load_fields="
@@ -68,7 +71,7 @@ public Operator get(DriverContext driverContext) {
                 parentTask,
                 maxOutstandingRequests,
                 inputChannel,
-                lookupService,
+                lookupService.apply(driverContext),
                 inputDataType,
                 lookupIndex,
                 matchField,

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

@@ -46,15 +46,15 @@ public class LookupFromIndexService extends AbstractLookupService<LookupFromInde
 
     public LookupFromIndexService(
         ClusterService clusterService,
-        CreateShardContext createShardContext,
+        LookupShardContextFactory lookupShardContextFactory,
         TransportService transportService,
         BigArrays bigArrays,
         BlockFactory blockFactory
     ) {
         super(
             LOOKUP_ACTION_NAME,
             clusterService,
-            createShardContext,
+            lookupShardContextFactory,
             transportService,
             bigArrays,
             blockFactory,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -589,7 +589,7 @@ private PhysicalOperation planLookupJoin(LookupJoinExec join, LocalExecutionPlan
                 parentTask,
                 context.queryPragmas().enrichMaxWorkers(),
                 matchFields.get(0).channel(),
-                lookupFromIndexService,
+                ctx -> lookupFromIndexService,
                 matchFields.get(0).type(),
                 indexName,
                 join.leftFields().get(0).name(),

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

@@ -108,19 +108,18 @@ public TransportEsqlQueryAction(
         exchangeService.registerTransportHandler(transportService);
         this.exchangeService = exchangeService;
         this.enrichPolicyResolver = new EnrichPolicyResolver(clusterService, transportService, planExecutor.indexResolver());
-        AbstractLookupService.CreateShardContext lookupCreateShardContext = AbstractLookupService.CreateShardContext.fromSearchService(
-            searchService
-        );
+        AbstractLookupService.LookupShardContextFactory lookupLookupShardContextFactory = AbstractLookupService.LookupShardContextFactory
+            .fromSearchService(searchService);
         this.enrichLookupService = new EnrichLookupService(
             clusterService,
-            lookupCreateShardContext,
+            lookupLookupShardContextFactory,
             transportService,
             bigArrays,
             blockFactory
         );
         this.lookupFromIndexService = new LookupFromIndexService(
             clusterService,
-            lookupCreateShardContext,
+            lookupLookupShardContextFactory,
             transportService,
             bigArrays,
             blockFactory