Return failures in response

Author: dnhatn

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -185,6 +185,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_THREAD_NAME_IN_DRIVER_PROFILE = def(9_027_0_00);
     public static final TransportVersion INFERENCE_CONTEXT = def(9_028_0_00);
     public static final TransportVersion ML_INFERENCE_DEEPSEEK = def(9_029_00_0);
+    public static final TransportVersion ESQL_FAILURE_FROM_REMOTE = def(9_030_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -13,6 +13,7 @@
 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;
@@ -136,4 +137,27 @@ 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

@@ -251,7 +251,7 @@ void runComputeOnRemoteCluster(
         try (var computeListener = new ComputeListener(transportService.getThreadPool(), cancelQueryOnFailure, listener.map(profiles -> {
             final TimeValue took = TimeValue.timeValueNanos(System.nanoTime() - startTimeInNanos);
             final ComputeResponse r = finalResponse.get();
-            return new ComputeResponse(profiles, took, r.totalShards, r.successfulShards, r.skippedShards, r.failedShards);
+            return new ComputeResponse(profiles, took, r.totalShards, r.successfulShards, r.skippedShards, r.failedShards, r.failures);
         }))) {
             var exchangeSource = new ExchangeSourceHandler(
                 configuration.pragmas().exchangeBufferSize(),

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

@@ -29,9 +29,10 @@ final class ComputeResponse extends TransportResponse {
     public final int successfulShards;
     public final int skippedShards;
     public final int failedShards;
+    public final List<Exception> failures;
 
     ComputeResponse(List<DriverProfile> profiles) {
-        this(profiles, null, null, null, null, null);
+        this(profiles, null, null, null, null, null, List.of());
     }
 
     ComputeResponse(
@@ -40,14 +41,16 @@ final class ComputeResponse extends TransportResponse {
         Integer totalShards,
         Integer successfulShards,
         Integer skippedShards,
-        Integer failedShards
+        Integer failedShards,
+        List<Exception> failures
     ) {
         this.profiles = profiles;
         this.took = took;
         this.totalShards = totalShards == null ? 0 : totalShards.intValue();
         this.successfulShards = successfulShards == null ? 0 : successfulShards.intValue();
         this.skippedShards = skippedShards == null ? 0 : skippedShards.intValue();
         this.failedShards = failedShards == null ? 0 : failedShards.intValue();
+        this.failures = failures;
     }
 
     ComputeResponse(StreamInput in) throws IOException {
@@ -74,6 +77,11 @@ final class ComputeResponse extends TransportResponse {
             this.skippedShards = 0;
             this.failedShards = 0;
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {
+            this.failures = in.readCollectionAsImmutableList(StreamInput::readException);
+        } else {
+            this.failures = List.of();
+        }
     }
 
     @Override
@@ -93,6 +101,9 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeVInt(skippedShards);
             out.writeVInt(failedShards);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_FAILURE_FROM_REMOTE)) {
+            out.writeCollection(failures, StreamOutput::writeException);
+        }
     }
 
     public List<DriverProfile> getProfiles() {
@@ -118,4 +129,8 @@ public int getSkippedShards() {
     public int getFailedShards() {
         return failedShards;
     }
+
+    public List<Exception> getFailures() {
+        return failures;
+    }
 }

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

@@ -115,7 +115,8 @@ final void startComputeOnDataNodes(
                     targetShards.totalShards(),
                     targetShards.totalShards() - shardFailures.size() - skippedShards.get(),
                     targetShards.skippedShards() + skippedShards.get(),
-                    shardFailures.size()
+                    shardFailures.size(),
+                    selectFailures()
                 );
             }))) {
                 for (TargetShard shard : targetShards.shards.values()) {
@@ -208,6 +209,18 @@ private void reportFailures(ComputeListener computeListener) {
         }
     }
 
+    private List<Exception> selectFailures() {
+        assert reportedFailure == false;
+        FailureCollector collector = new FailureCollector();
+        Set<Exception> seen = Collections.newSetFromMap(new IdentityHashMap<>());
+        for (ShardFailure e : shardFailures.values()) {
+            if (seen.add(e.failure)) {
+                collector.unwrapAndCollect(e.failure);
+            }
+        }
+        return collector.getFailures();
+    }
+
     private void sendOneNodeRequest(TargetShards targetShards, ComputeListener computeListener, NodeRequest request) {
         final ActionListener<List<DriverProfile>> listener = computeListener.acquireCompute();
         sendRequest(request.node, request.shardIds, request.aliasFilters, new NodeListener() {