Return failures in partial response

dnhatn · dnhatn · commit b71ca6b65280 · 2025-03-14T16:02:17.000-07:00
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FailureCollector.java b/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;
-        }
-    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ClusterComputeHandler.java b/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);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeResponse.java b/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;
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java b/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());
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java b/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()
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java b/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) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSenderTests.java b/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

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@`
`8`	`8`	`package org.elasticsearch.xpack.esql.plugin;`
`9`	`9`
`10`	`10`	`import org.elasticsearch.TransportVersions;`
	`11`	`+import org.elasticsearch.action.search.ShardSearchFailure;`
`11`	`12`	`import org.elasticsearch.common.io.stream.StreamInput;`
`12`	`13`	`import org.elasticsearch.common.io.stream.StreamOutput;`
`13`	`14`	`import org.elasticsearch.compute.operator.DriverProfile;`
`@@ -29,7 +30,7 @@ final class ComputeResponse extends TransportResponse {`
`29`	`30`	`public final int successfulShards;`
`30`	`31`	`public final int skippedShards;`
`31`	`32`	`public final int failedShards;`
`32`		`- public final List<Exception> failures;`
	`33`	`+ public final List<ShardSearchFailure> failures;`
`33`	`34`
`34`	`35`	`ComputeResponse(List<DriverProfile> profiles) {`
`35`	`36`	`this(profiles, null, null, null, null, null, List.of());`
`@@ -42,7 +43,7 @@ final class ComputeResponse extends TransportResponse {`
`42`	`43`	`Integer successfulShards,`
`43`	`44`	`Integer skippedShards,`
`44`	`45`	`Integer failedShards,`
`45`		`- List<Exception> failures`
	`46`	`+ List<ShardSearchFailure> failures`
`46`	`47`	`) {`
`47`	`48`	`this.profiles = profiles;`
`48`	`49`	`this.took = took;`
`@@ -78,7 +79,7 @@ final class ComputeResponse extends TransportResponse {`
`78`	`79`	`this.failedShards = 0;`
`79`	`80`	`}`
`80`	`81`	`if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {`
`81`		`- this.failures = in.readCollectionAsImmutableList(StreamInput::readException);`
	`82`	`+ this.failures = in.readCollectionAsImmutableList(ShardSearchFailure::readShardSearchFailure);`
`82`	`83`	`} else {`
`83`	`84`	`this.failures = List.of();`
`84`	`85`	`}`
`@@ -102,7 +103,7 @@ public void writeTo(StreamOutput out) throws IOException {`
`102`	`103`	`out.writeVInt(failedShards);`
`103`	`104`	`}`
`104`	`105`	`if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {`
`105`		`- out.writeCollection(failures, StreamOutput::writeException);`
	`106`	`+ out.writeCollection(failures, (o, v) -> v.writeTo(o));`
`106`	`107`	`}`
`107`	`108`	`}`
`108`	`109`
`@@ -130,7 +131,7 @@ public int getFailedShards() {`
`130`	`131`	`return failedShards;`
`131`	`132`	`}`
`132`	`133`
`133`		`- public List<Exception> getFailures() {`
	`134`	`+ public List<ShardSearchFailure> getFailures() {`
`134`	`135`	`return failures;`
`135`	`136`	`}`
`136`	`137`	`}`