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() {