Make INLINESTATS (and subplans) work with CCS

Author: smalyshev

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

@@ -83,6 +83,9 @@ public class EsqlExecutionInfo implements ChunkedToXContentObject, Writeable {
     private transient TimeSpan planningTimeSpan; // time elapsed since start of query to calling ComputeService.execute
     private TimeValue overallTook;
 
+    // Are we doing subplans? No need to serialize this because it is only relevant for the coordinator node.
+    private transient boolean inSubplan = false;
+
     // This is only used is tests.
     public EsqlExecutionInfo(boolean includeCCSMetadata) {
         this(Predicates.always(), includeCCSMetadata);  // default all clusters to being skippable on failure
@@ -142,6 +145,7 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeOptionalWriteable(overallTimeSpan);
             out.writeOptionalWriteable(planningTimeSpan);
         }
+        assert inSubplan == false : "Should not be serializing execution info while in subplans";
     }
 
     public boolean includeCCSMetadata() {
@@ -168,8 +172,10 @@ public TimeValue planningTookTime() {
      */
     public void markEndQuery() {
         assert relativeStart != null : "Relative start time must be set when markEndQuery is called";
-        overallTimeSpan = relativeStart.stop();
-        overallTook = overallTimeSpan.toTimeValue();
+        if (isMainPlan()) {
+            overallTimeSpan = relativeStart.stop();
+            overallTook = overallTimeSpan.toTimeValue();
+        }
     }
 
     // for testing only - use markEndQuery in production code
@@ -358,6 +364,18 @@ public void clusterInfoInitializing(boolean clusterInfoInitializing) {
         this.clusterInfoInitializing = clusterInfoInitializing;
     }
 
+    public boolean isMainPlan() {
+        return inSubplan == false;
+    }
+
+    public void startSubPlan() {
+        this.inSubplan = true;
+    }
+
+    public void finishSubPlan() {
+        this.inSubplan = false;
+    }
+
     /**
      * Represents the search metadata about a particular cluster involved in a cross-cluster search.
      * The Cluster object can represent either the local cluster or a remote cluster.

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

@@ -158,14 +158,15 @@ private void updateExecutionInfo(EsqlExecutionInfo executionInfo, String cluster
                 .setSuccessfulShards(resp.getSuccessfulShards())
                 .setSkippedShards(resp.getSkippedShards())
                 .setFailedShards(resp.getFailedShards());
+            // TODO: correctly aggregate took time from subplans
             if (resp.getTook() != null) {
                 builder.setTook(TimeValue.timeValueNanos(executionInfo.planningTookTime().nanos() + resp.getTook().nanos()));
             } else {
                 // if the cluster is an older version and does not send back took time, then calculate it here on the coordinator
                 // and leave shard info unset, so it is not shown in the CCS metadata section of the JSON response
                 builder.setTook(executionInfo.tookSoFar());
             }
-            if (v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) {
+            if (executionInfo.isMainPlan() && v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) {
                 builder.addFailures(resp.failures);
                 if (executionInfo.isStopped() || resp.failedShards > 0 || resp.failures.isEmpty() == false) {
                     builder.setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL);

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

@@ -384,7 +384,7 @@ public void executePlan(
                 cancelQueryOnFailure,
                 listener.delegateFailureAndWrap((l, completionInfo) -> {
                     failIfAllShardsFailed(execInfo, collectedPages);
-                    execInfo.markEndQuery();  // TODO: revisit this time recording model as part of INLINESTATS improvements
+                    execInfo.markEndQuery();
                     l.onResponse(new Result(outputAttributes, collectedPages, completionInfo, execInfo));
                 })
             )
@@ -401,7 +401,7 @@ public void executePlan(
                                 execInfo.swapCluster(LOCAL_CLUSTER, (k, v) -> {
                                     var tookTime = execInfo.tookSoFar();
                                     var builder = new EsqlExecutionInfo.Cluster.Builder(v).setTook(tookTime);
-                                    if (v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) {
+                                    if (execInfo.isMainPlan() && v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) {
                                         final Integer failedShards = execInfo.getCluster(LOCAL_CLUSTER).getFailedShards();
                                         // Set the local cluster status (including the final driver) to partial if the query was stopped
                                         // or encountered resolution or execution failures.
@@ -533,8 +533,8 @@ private static void updateShardCountForCoordinatorOnlyQuery(EsqlExecutionInfo ex
 
     // For queries like: FROM logs* | LIMIT 0 (including cross-cluster LIMIT 0 queries)
     private static void updateExecutionInfoAfterCoordinatorOnlyQuery(EsqlExecutionInfo execInfo) {
-        execInfo.markEndQuery();  // TODO: revisit this time recording model as part of INLINESTATS improvements
-        if (execInfo.isCrossClusterSearch()) {
+        execInfo.markEndQuery();
+        if (execInfo.isCrossClusterSearch() && execInfo.isMainPlan()) {
             assert execInfo.planningTookTime() != null : "Planning took time should be set on EsqlExecutionInfo but is null";
             for (String clusterAlias : execInfo.clusterAliases()) {
                 execInfo.swapCluster(clusterAlias, (k, v) -> {

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

@@ -121,6 +121,7 @@ static boolean returnSuccessWithEmptyResult(EsqlExecutionInfo executionInfo, Exc
     }
 
     static void updateExecutionInfoToReturnEmptyResult(EsqlExecutionInfo executionInfo, Exception e) {
+        // This applies even for subplans - if we had an error and have to skip a cluster, then it will remain skipped.
         executionInfo.markEndQuery();
         Exception exceptionForResponse;
         if (e instanceof ConnectTransportException) {

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

@@ -262,7 +262,13 @@ private void executeSubPlan(
         // Create a physical plan out of the logical sub-plan
         var physicalSubPlan = logicalPlanToPhysicalPlan(subPlans.stubReplacedSubPlan(), request);
 
-        runner.run(physicalSubPlan, listener.delegateFailureAndWrap((next, result) -> {
+        executionInfo.startSubPlan();
+        var listenerForSublan = listener.delegateResponse((l, e) -> {
+            executionInfo.finishSubPlan();
+            l.onFailure(e);
+        });
+
+        runner.run(physicalSubPlan, listenerForSublan.delegateFailureAndWrap((next, result) -> {
             try {
                 // Translate the subquery into a separate, coordinator based plan and the results 'broadcasted' as a local relation
                 completionInfoAccumulator.accumulate(result.completionInfo());
@@ -282,6 +288,8 @@ private void executeSubPlan(
                 var newSubPlan = firstSubPlan(newLogicalPlan);
 
                 if (newSubPlan == null) {// run the final "main" plan
+                    // TODO: failures!
+                    executionInfo.finishSubPlan();
                     LOGGER.debug("Executing final plan:\n{}", newLogicalPlan);
                     var newPhysicalPlan = logicalPlanToPhysicalPlan(newLogicalPlan, request);
                     runner.run(newPhysicalPlan, next.delegateFailureAndWrap((finalListener, finalResult) -> {
@@ -291,7 +299,15 @@ private void executeSubPlan(
                         );
                     }));
                 } else {// continue executing the subplans
-                    executeSubPlan(completionInfoAccumulator, newLogicalPlan, newSubPlan, executionInfo, runner, request, listener);
+                    executeSubPlan(
+                        completionInfoAccumulator,
+                        newLogicalPlan,
+                        newSubPlan,
+                        executionInfo,
+                        runner,
+                        request,
+                        listenerForSublan
+                    );
                 }
             } finally {
                 Releasables.closeExpectNoException(Releasables.wrap(Iterators.map(result.pages().iterator(), p -> p::releaseBlocks)));