Make analyzer async.

Author: afoucret

benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/QueryPlanningBenchmark.java

@@ -9,6 +9,9 @@
 
 package org.elasticsearch.benchmark._nightly.esql;
 
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.IndexMode;
@@ -49,6 +52,7 @@
 import java.util.LinkedHashMap;
 import java.util.Locale;
 import java.util.Map;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
 import static java.util.Collections.emptyMap;
@@ -115,15 +119,22 @@ public void setup() {
         defaultOptimizer = new LogicalPlanOptimizer(new LogicalOptimizerContext(config, FoldContext.small()));
     }
 
-    private LogicalPlan plan(EsqlParser parser, Analyzer analyzer, LogicalPlanOptimizer optimizer, String query) {
+    private void plan(
+        EsqlParser parser,
+        Analyzer analyzer,
+        LogicalPlanOptimizer optimizer,
+        String query,
+        ActionListener<LogicalPlan> listener
+    ) {
         var parsed = parser.createStatement(query, new QueryParams(), telemetry);
-        var analyzed = analyzer.analyze(parsed);
-        var optimized = optimizer.optimize(analyzed);
-        return optimized;
+        SubscribableListener.<LogicalPlan>newForked(analyzedPlanListener -> analyzer.analyze(parsed, analyzedPlanListener))
+            .addListener(listener.map(optimizer::optimize));
     }
 
     @Benchmark
-    public void manyFields(Blackhole blackhole) {
-        blackhole.consume(plan(defaultParser, manyFieldsAnalyzer, defaultOptimizer, "FROM test | LIMIT 10"));
+    public void manyFields(Blackhole blackhole) throws ExecutionException, InterruptedException {
+        PlainActionFuture<LogicalPlan> optimizedPlanFuture = new PlainActionFuture<>();
+        plan(defaultParser, manyFieldsAnalyzer, defaultOptimizer, "FROM test | LIMIT 10", optimizedPlanFuture);
+        blackhole.consume(optimizedPlanFuture.get());
     }
 }

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.analysis;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.logging.HeaderWarning;
 import org.elasticsearch.common.logging.LoggerMessageFormat;
 import org.elasticsearch.common.lucene.BytesRefs;
@@ -194,9 +195,11 @@ public Analyzer(AnalyzerContext context, Verifier verifier) {
         this.verifier = verifier;
     }
 
-    public LogicalPlan analyze(LogicalPlan plan) {
-        BitSet partialMetrics = new BitSet(FeatureMetric.values().length);
-        return verify(execute(plan), gatherPreAnalysisMetrics(plan, partialMetrics));
+    public void analyze(LogicalPlan plan, ActionListener<LogicalPlan> listener) {
+        execute(plan, listener.map(analzedPlan -> {
+            BitSet partialMetrics = new BitSet(FeatureMetric.values().length);
+            return this.verify(analzedPlan, gatherPreAnalysisMetrics(plan, partialMetrics));
+        }));
     }
 
     public LogicalPlan verify(LogicalPlan plan, BitSet partialMetrics) {

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

@@ -108,7 +108,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Function;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
@@ -342,14 +342,22 @@ public void analyzedPlan(
             return;
         }
 
-        Function<PreAnalysisResult, LogicalPlan> analyzeAction = (l) -> {
+        BiConsumer<PreAnalysisResult, ActionListener<LogicalPlan>> analyzeAction = (preAnalysisResult, listener) -> {
             Analyzer analyzer = new Analyzer(
-                new AnalyzerContext(configuration, functionRegistry, l.indices, l.lookupIndices, l.enrichResolution, l.inferenceResolution),
+                new AnalyzerContext(
+                    configuration,
+                    functionRegistry,
+                    preAnalysisResult.indices,
+                    preAnalysisResult.lookupIndices,
+                    preAnalysisResult.enrichResolution,
+                    preAnalysisResult.inferenceResolution
+                ),
                 verifier
             );
-            LogicalPlan plan = analyzer.analyze(parsed);
-            plan.setAnalyzed();
-            return plan;
+            analyzer.analyze(parsed, listener.map(analyzedPlan -> {
+                analyzedPlan.setAnalyzed();
+                return analyzedPlan;
+            }));
         };
         // Capture configured remotes list to ensure consistency throughout the session
         configuredClusters = Set.copyOf(indicesExpressionGrouper.getConfiguredClusters());
@@ -404,13 +412,14 @@ public void analyzedPlan(
                 // do we want to declare it successful or skipped? For now, unavailability takes precedence.
                 EsqlCCSUtils.updateExecutionInfoWithUnavailableClusters(executionInfo, result.indices.unavailableClusters());
                 EsqlCCSUtils.updateExecutionInfoWithClustersWithNoMatchingIndices(executionInfo, result.indices, null);
-                plan = analyzeAction.apply(result);
+                analyzeAction.accept(result, l.delegateFailureAndWrap((analyzedPlanListener, analyzedPlan) -> {
+                    LOGGER.debug("Analyzed plan (second attempt, without filter):\n{}", analyzedPlan);
+                    analyzedPlanListener.onResponse(analyzedPlan);
+                }));
             } catch (Exception e) {
                 l.onFailure(e);
                 return;
             }
-            LOGGER.debug("Analyzed plan (second attempt, without filter):\n{}", plan);
-            l.onResponse(plan);
         }).addListener(logicalPlanListener);
     }
 
@@ -528,7 +537,7 @@ private boolean allCCSClustersSkipped(
     }
 
     private static void analyzeAndMaybeRetry(
-        Function<PreAnalysisResult, LogicalPlan> analyzeAction,
+        BiConsumer<PreAnalysisResult, ActionListener<LogicalPlan>> analyzeAction,
         QueryBuilder requestFilter,
         PreAnalysisResult result,
         EsqlExecutionInfo executionInfo,
@@ -540,14 +549,7 @@ private static void analyzeAndMaybeRetry(
         var attemptMessage = requestFilter == null ? "the only" : "first";
         LOGGER.debug("Analyzing the plan ({} attempt, {} filter)", attemptMessage, filterPresentMessage);
 
-        try {
-            if (result.indices.isValid() || requestFilter != null) {
-                // We won't run this check with no filter and no valid indices since this may lead to false positive - missing index report
-                // when the resolution result is not valid for a different reason.
-                EsqlCCSUtils.updateExecutionInfoWithClustersWithNoMatchingIndices(executionInfo, result.indices, requestFilter);
-            }
-            plan = analyzeAction.apply(result);
-        } catch (Exception e) {
+        logicalPlanListener = logicalPlanListener.delegateResponse((errorListener, e) -> {
             if (e instanceof VerificationException ve) {
                 LOGGER.debug(
                     "Analyzing the plan ({} attempt, {} filter) failed with {}",
@@ -557,21 +559,31 @@ private static void analyzeAndMaybeRetry(
                 );
                 if (requestFilter == null) {
                     // if the initial request didn't have a filter, then just pass the exception back to the user
-                    logicalPlanListener.onFailure(ve);
+                    errorListener.onFailure(ve);
                 } else {
                     // interested only in a VerificationException, but this time we are taking out the index filter
                     // to try and make the index resolution work without any index filtering. In the next step... to be continued
                     l.onResponse(result);
                 }
             } else {
                 // if the query failed with any other type of exception, then just pass the exception back to the user
-                logicalPlanListener.onFailure(e);
+                errorListener.onFailure(e);
             }
-            return;
+        });
+
+        try {
+            if (result.indices.isValid() || requestFilter != null) {
+                // We won't run this check with no filter and no valid indices since this may lead to false positive - missing index report
+                // when the resolution result is not valid for a different reason.
+                EsqlCCSUtils.updateExecutionInfoWithClustersWithNoMatchingIndices(executionInfo, result.indices, requestFilter);
+            }
+            analyzeAction.accept(result, ActionListener.runBefore(logicalPlanListener, () -> {
+                LOGGER.debug("Analyzed plan ({} attempt, {} filter):\n{}", attemptMessage, filterPresentMessage, plan);
+                // the analysis succeeded from the first attempt, irrespective if it had a filter or not, just continue with the planning
+            }));
+        } catch (Exception e) {
+            l.onFailure(e);
         }
-        LOGGER.debug("Analyzed plan ({} attempt, {} filter):\n{}", attemptMessage, filterPresentMessage, plan);
-        // the analysis succeeded from the first attempt, irrespective if it had a filter or not, just continue with the planning
-        logicalPlanListener.onResponse(plan);
     }
 
     private static void resolveFieldNames(LogicalPlan parsed, EnrichResolution enrichResolution, ActionListener<PreAnalysisResult> l) {