Return failures in partial response

Author: dnhatn

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FailureCollector.java

@@ -13,7 +13,6 @@
 import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.transport.TransportException;
 
-import java.util.ArrayList;
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
@@ -137,27 +136,4 @@ private Exception buildFailure() {
         assert first != null;
         return first;
     }
-
-    public List<Exception> getFailures() {
-        if (hasFailure == false) {
-            return List.of();
-        }
-        synchronized (this) {
-            List<Exception> failures = new ArrayList<>();
-            for (Category category : List.of(Category.CLIENT, Category.SERVER, Category.SHARD_UNAVAILABLE)) {
-                for (Exception e : categories.get(category)) {
-                    if (failures.size() <= maxExceptions) {
-                        failures.add(e);
-                    }
-                }
-            }
-            if (failures.isEmpty()) {
-                Exception any = categories.get(Category.CANCELLATION).poll();
-                if (any != null) {
-                    failures.add(any);
-                }
-            }
-            return failures;
-        }
-    }
 }

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

@@ -163,7 +163,8 @@ private void updateExecutionInfo(EsqlExecutionInfo executionInfo, String cluster
                 builder.setTook(executionInfo.tookSoFar());
             }
             if (v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) {
-                if (executionInfo.isStopped() || resp.failedShards > 0) {
+                builder.setFailures(resp.failures);
+                if (executionInfo.isStopped() || resp.failedShards > 0 || resp.failures.isEmpty() == false) {
                     builder.setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL);
                 } else {
                     builder.setStatus(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL);

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.plugin;
 
 import org.elasticsearch.TransportVersions;
+import org.elasticsearch.action.search.ShardSearchFailure;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.operator.DriverProfile;
@@ -29,7 +30,7 @@ final class ComputeResponse extends TransportResponse {
     public final int successfulShards;
     public final int skippedShards;
     public final int failedShards;
-    public final List<Exception> failures;
+    public final List<ShardSearchFailure> failures;
 
     ComputeResponse(List<DriverProfile> profiles) {
         this(profiles, null, null, null, null, null, List.of());
@@ -42,7 +43,7 @@ final class ComputeResponse extends TransportResponse {
         Integer successfulShards,
         Integer skippedShards,
         Integer failedShards,
-        List<Exception> failures
+        List<ShardSearchFailure> failures
     ) {
         this.profiles = profiles;
         this.took = took;
@@ -78,7 +79,7 @@ final class ComputeResponse extends TransportResponse {
             this.failedShards = 0;
         }
         if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {
-            this.failures = in.readCollectionAsImmutableList(StreamInput::readException);
+            this.failures = in.readCollectionAsImmutableList(ShardSearchFailure::readShardSearchFailure);
         } else {
             this.failures = List.of();
         }
@@ -102,7 +103,7 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeVInt(failedShards);
         }
         if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {
-            out.writeCollection(failures, StreamOutput::writeException);
+            out.writeCollection(failures, (o, v) -> v.writeTo(o));
         }
     }
 
@@ -130,7 +131,7 @@ public int getFailedShards() {
         return failedShards;
     }
 
-    public List<Exception> getFailures() {
+    public List<ShardSearchFailure> getFailures() {
         return failures;
     }
 }

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

@@ -268,6 +268,7 @@ public void execute(
                                         .setSuccessfulShards(r.getSuccessfulShards())
                                         .setSkippedShards(r.getSkippedShards())
                                         .setFailedShards(r.getFailedShards())
+                                        .setFailures(r.failures)
                                         .build()
                                 );
                                 dataNodesListener.onResponse(r.getProfiles());

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

@@ -101,6 +101,7 @@ void startComputeOnDataNodes(
         new DataNodeRequestSender(
             transportService,
             esqlExecutor,
+            clusterAlias,
             parentTask,
             configuration.allowPartialResults(),
             configuration.pragmas().maxConcurrentNodesPerCluster()

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

@@ -15,6 +15,7 @@
 import org.elasticsearch.action.search.SearchShardsGroup;
 import org.elasticsearch.action.search.SearchShardsRequest;
 import org.elasticsearch.action.search.SearchShardsResponse;
+import org.elasticsearch.action.search.ShardSearchFailure;
 import org.elasticsearch.action.support.TransportActions;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
@@ -26,6 +27,7 @@
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
@@ -73,6 +75,7 @@ abstract class DataNodeRequestSender {
 
     private final TransportService transportService;
     private final Executor esqlExecutor;
+    private final String clusterAlias;
     private final CancellableTask rootTask;
     private final boolean allowPartialResults;
     private final Semaphore concurrentRequests;
@@ -87,12 +90,14 @@ abstract class DataNodeRequestSender {
     DataNodeRequestSender(
         TransportService transportService,
         Executor esqlExecutor,
+        String clusterAlias,
         CancellableTask rootTask,
         boolean allowPartialResults,
         int concurrentRequests
     ) {
         this.transportService = transportService;
         this.esqlExecutor = esqlExecutor;
+        this.clusterAlias = clusterAlias;
         this.rootTask = rootTask;
         this.allowPartialResults = allowPartialResults;
         this.concurrentRequests = concurrentRequests > 0 ? new Semaphore(concurrentRequests) : null;
@@ -209,16 +214,25 @@ private void reportFailures(ComputeListener computeListener) {
         }
     }
 
-    private List<Exception> selectFailures() {
+    private List<ShardSearchFailure> selectFailures() {
         assert reportedFailure == false;
-        FailureCollector collector = new FailureCollector();
+        List<ShardSearchFailure> failures = new ArrayList<>();
         Set<Exception> seen = Collections.newSetFromMap(new IdentityHashMap<>());
-        for (ShardFailure e : shardFailures.values()) {
-            if (seen.add(e.failure)) {
-                collector.unwrapAndCollect(e.failure);
+        for (Map.Entry<ShardId, ShardFailure> e : shardFailures.entrySet()) {
+            ShardFailure failure = e.getValue();
+            if (ExceptionsHelper.unwrap(failure.failure(), TaskCancelledException.class) != null) {
+                continue;
             }
+            if (seen.add(failure.failure) && failures.size() < 5) {
+                failures.add(new ShardSearchFailure(failure.failure, new SearchShardTarget(null, e.getKey(), clusterAlias)));
+            }
+        }
+        // pick any cancellation
+        if (failures.isEmpty() && shardFailures.isEmpty() == false) {
+            ShardFailure any = shardFailures.values().iterator().next();
+            failures.add(new ShardSearchFailure(any.failure));
         }
-        return collector.getFailures();
+        return failures;
     }
 
     private void sendOneNodeRequest(TargetShards targetShards, ComputeListener computeListener, NodeRequest request) {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSenderTests.java

@@ -453,6 +453,7 @@ PlainActionFuture<ComputeResponse> sendRequests(
         DataNodeRequestSender requestSender = new DataNodeRequestSender(
             transportService,
             executor,
+            "",
             task,
             allowPartialResults,
             concurrentRequests