Add handling remote sink failures

Author: smalyshev

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

@@ -42,6 +42,7 @@
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 /**
  * {@link ExchangeService} is responsible for exchanging pages between exchange sinks and sources on the same or different nodes.
@@ -277,9 +278,16 @@ protected void doRun() {
      * @param exchangeId       the exchange ID
      * @param transportService the transport service
      * @param conn             the connection to the remote node where the remote exchange sink is located
+     * @param failureCollector if not null, the failures will be sent to this consumer and the sink will be marked as finished
      */
-    public RemoteSink newRemoteSink(Task parentTask, String exchangeId, TransportService transportService, Transport.Connection conn) {
-        return new TransportRemoteSink(transportService, blockFactory, conn, parentTask, exchangeId, executor);
+    public RemoteSink newRemoteSink(
+        Task parentTask,
+        String exchangeId,
+        TransportService transportService,
+        Transport.Connection conn,
+        Consumer<Exception> failureCollector
+    ) {
+        return new TransportRemoteSink(transportService, blockFactory, conn, parentTask, exchangeId, executor, failureCollector);
     }
 
     static final class TransportRemoteSink implements RemoteSink {
@@ -289,6 +297,7 @@ static final class TransportRemoteSink implements RemoteSink {
         final Task parentTask;
         final String exchangeId;
         final Executor responseExecutor;
+        final Consumer<Exception> failureCollector;
 
         final AtomicLong estimatedPageSizeInBytes = new AtomicLong(0L);
 
@@ -298,14 +307,16 @@ static final class TransportRemoteSink implements RemoteSink {
             Transport.Connection connection,
             Task parentTask,
             String exchangeId,
-            Executor responseExecutor
+            Executor responseExecutor,
+            Consumer<Exception> failureCollector
         ) {
             this.transportService = transportService;
             this.blockFactory = blockFactory;
             this.connection = connection;
             this.parentTask = parentTask;
             this.exchangeId = exchangeId;
             this.responseExecutor = responseExecutor;
+            this.failureCollector = failureCollector;
         }
 
         @Override
@@ -316,6 +327,12 @@ public void fetchPageAsync(boolean allSourcesFinished, ActionListener<ExchangeRe
                 blockFactory.breaker().addEstimateBytesAndMaybeBreak(reservedBytes, "fetch page");
                 listener = ActionListener.runAfter(listener, () -> blockFactory.breaker().addWithoutBreaking(-reservedBytes));
             }
+            if (failureCollector != null) {
+                listener = listener.delegateResponse((l, ex) -> {
+                    failureCollector.accept(ex);
+                    l.onResponse(new ExchangeResponse(blockFactory, null, true));
+                });
+            }
             transportService.sendChildRequest(
                 connection,
                 EXCHANGE_ACTION_NAME,

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java

@@ -383,7 +383,7 @@ public void testConcurrentWithTransportActions() {
             sourceHandler.addCompletionListener(sourceCompletionFuture);
             ExchangeSinkHandler sinkHandler = exchange1.createSinkHandler(exchangeId, randomExchangeBuffer());
             Transport.Connection connection = node0.getConnection(node1.getLocalNode());
-            sourceHandler.addRemoteSink(exchange0.newRemoteSink(task, exchangeId, node0, connection), randomIntBetween(1, 5));
+            sourceHandler.addRemoteSink(exchange0.newRemoteSink(task, exchangeId, node0, connection, null), randomIntBetween(1, 5));
             final int maxInputSeqNo = rarely() ? -1 : randomIntBetween(0, 50_000);
             final int maxOutputSeqNo = rarely() ? -1 : randomIntBetween(0, 50_000);
             runConcurrentTest(maxInputSeqNo, maxOutputSeqNo, sourceHandler::createExchangeSource, sinkHandler::createExchangeSink);
@@ -442,7 +442,7 @@ public void sendResponse(TransportResponse transportResponse) {
             sourceHandler.addCompletionListener(sourceCompletionFuture);
             ExchangeSinkHandler sinkHandler = exchange1.createSinkHandler(exchangeId, randomIntBetween(1, 128));
             Transport.Connection connection = node0.getConnection(node1.getLocalNode());
-            sourceHandler.addRemoteSink(exchange0.newRemoteSink(task, exchangeId, node0, connection), randomIntBetween(1, 5));
+            sourceHandler.addRemoteSink(exchange0.newRemoteSink(task, exchangeId, node0, connection, null), randomIntBetween(1, 5));
             Exception err = expectThrows(
                 Exception.class,
                 () -> runConcurrentTest(maxSeqNo, maxSeqNo, sourceHandler::createExchangeSource, sinkHandler::createExchangeSink)

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

@@ -202,6 +202,9 @@ void markAsPartial(String computeClusterAlias, Exception e) {
         esqlExecutionInfo.swapCluster(computeClusterAlias, (k, v) -> {
             assert v.getStatus() != EsqlExecutionInfo.Cluster.Status.SKIPPED
                 : "We shouldn't be running compute on a cluster that's already marked as skipped";
+            if (v.getStatus() == EsqlExecutionInfo.Cluster.Status.PARTIAL) {
+                return v;
+            }
             return new EsqlExecutionInfo.Cluster.Builder(v).setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL)
                 .setFailures(List.of(new ShardSearchFailure(e)))
                 .build();

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

@@ -335,7 +335,7 @@ private void startComputeOnDataNodes(
                         queryPragmas.exchangeBufferSize(),
                         esqlExecutor,
                         refs.acquire().delegateFailureAndWrap((l, unused) -> {
-                            var remoteSink = exchangeService.newRemoteSink(parentTask, sessionId, transportService, node.connection);
+                            var remoteSink = exchangeService.newRemoteSink(parentTask, sessionId, transportService, node.connection, null);
                             exchangeSource.addRemoteSink(remoteSink, queryPragmas.concurrentExchangeClients());
                             ActionListener<ComputeResponse> computeResponseListener = computeListener.acquireCompute(clusterAlias);
                             var dataNodeListener = ActionListener.runBefore(computeResponseListener, () -> l.onResponse(null));
@@ -377,7 +377,7 @@ private void startComputeOnRemoteClusters(
         try (RefCountingListener refs = new RefCountingListener(linkExchangeListeners)) {
             for (RemoteCluster cluster : clusters) {
                 final String clusterAlias = cluster.clusterAlias();
-                final boolean shouldSkipOnFailure = computeListener.shouldIgnoreRemoteErrors(clusterAlias);
+                final boolean suppressRemoteFailure = computeListener.shouldIgnoreRemoteErrors(clusterAlias);
                 final var exchangeListener = refs.acquire();
                 ExchangeService.openExchange(
                     transportService,
@@ -386,7 +386,13 @@ private void startComputeOnRemoteClusters(
                     queryPragmas.exchangeBufferSize(),
                     esqlExecutor,
                     ActionListener.wrap(unused -> {
-                        var remoteSink = exchangeService.newRemoteSink(rootTask, sessionId, transportService, cluster.connection);
+                        var remoteSink = exchangeService.newRemoteSink(
+                            rootTask,
+                            sessionId,
+                            transportService,
+                            cluster.connection,
+                            suppressRemoteFailure ? ex -> computeListener.markAsPartial(clusterAlias, ex) : null
+                        );
                         exchangeSource.addRemoteSink(remoteSink, queryPragmas.concurrentExchangeClients());
                         var remotePlan = new RemoteClusterPlan(plan, cluster.concreteIndices, cluster.originalIndices);
                         var clusterRequest = new ClusterComputeRequest(clusterAlias, sessionId, configuration, remotePlan);
@@ -403,7 +409,7 @@ private void startComputeOnRemoteClusters(
                             new ActionListenerResponseHandler<>(clusterListener, ComputeResponse::new, esqlExecutor)
                         );
                     }, e -> {
-                        if (shouldSkipOnFailure) {
+                        if (suppressRemoteFailure) {
                             // TODO: drop this in final patch
                             LOGGER.error("Marking failed cluster {} as partial: {}", clusterAlias, e);
                             computeListener.markAsPartial(clusterAlias, e);