Cover more cases for skip

Author: smalyshev

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

@@ -189,23 +189,33 @@ ActionListener<Void> acquireAvoid() {
         });
     }
 
+    boolean shouldIgnoreRemoteErrors(@Nullable String computeClusterAlias) {
+        return computeClusterAlias != null
+            && esqlExecutionInfo.isCrossClusterSearch()
+            && computeClusterAlias.equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY) == false
+            && runningOnRemoteCluster() == false
+            && esqlExecutionInfo.isSkipUnavailable(computeClusterAlias);
+    }
+
+    void markAsPartial(String computeClusterAlias, Exception e) {
+        // We use PARTIAL here because we can not know whether the cluster have already sent any data.
+        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";
+            return new EsqlExecutionInfo.Cluster.Builder(v).setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL)
+                .setFailures(List.of(new ShardSearchFailure(e)))
+                .build();
+        });
+    }
+
     ActionListener<Void> acquireSkipUnavailable(@Nullable String computeClusterAlias) {
-        if (computeClusterAlias == null
-            || esqlExecutionInfo.isCrossClusterSearch() == false
-            || runningOnRemoteCluster()
-            || computeClusterAlias.equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)
-            || esqlExecutionInfo.isSkipUnavailable(computeClusterAlias) == false) {
+        if (shouldIgnoreRemoteErrors(computeClusterAlias) == false) {
             return acquireAvoid();
         }
         return refs.acquire().delegateResponse((l, e) -> {
-            LOGGER.error("Skipping unavailable cluster {} in ESQL query: {}", computeClusterAlias, 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";
-                return new EsqlExecutionInfo.Cluster.Builder(v).setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL)
-                    .setFailures(List.of(new ShardSearchFailure(e)))
-                    .build();
-            });
+            // TODO: drop this in final patch
+            LOGGER.error("Marking failed cluster {} as partial: {}", computeClusterAlias, e);
+            markAsPartial(computeClusterAlias, e);
             l.onResponse(null);
         });
     }

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

@@ -380,20 +380,23 @@ private void startComputeOnRemoteClusters(
         var linkExchangeListeners = ActionListener.releaseAfter(computeListener.acquireAvoid(), exchangeSource.addEmptySink());
         try (RefCountingListener refs = new RefCountingListener(linkExchangeListeners)) {
             for (RemoteCluster cluster : clusters) {
+                final String clusterAlias = cluster.clusterAlias();
+                final boolean shouldSkipOnFailure = computeListener.shouldIgnoreRemoteErrors(clusterAlias);
+                final var exchangeListener = refs.acquire();
                 ExchangeService.openExchange(
                     transportService,
                     cluster.connection,
                     sessionId,
                     queryPragmas.exchangeBufferSize(),
                     esqlExecutor,
-                    refs.acquire().delegateFailureAndWrap((l, unused) -> {
+                    ActionListener.wrap(unused -> {
                         var remoteSink = exchangeService.newRemoteSink(rootTask, sessionId, transportService, cluster.connection);
                         exchangeSource.addRemoteSink(remoteSink, queryPragmas.concurrentExchangeClients());
                         var remotePlan = new RemoteClusterPlan(plan, cluster.concreteIndices, cluster.originalIndices);
-                        var clusterRequest = new ClusterComputeRequest(cluster.clusterAlias, sessionId, configuration, remotePlan);
+                        var clusterRequest = new ClusterComputeRequest(clusterAlias, sessionId, configuration, remotePlan);
                         var clusterListener = ActionListener.runBefore(
-                            computeListener.acquireCompute(cluster.clusterAlias()),
-                            () -> l.onResponse(null)
+                            computeListener.acquireCompute(clusterAlias),
+                            () -> exchangeListener.onResponse(null)
                         );
                         transportService.sendChildRequest(
                             cluster.connection,
@@ -403,6 +406,15 @@ private void startComputeOnRemoteClusters(
                             TransportRequestOptions.EMPTY,
                             new ActionListenerResponseHandler<>(clusterListener, ComputeResponse::new, esqlExecutor)
                         );
+                    }, e -> {
+                        if (shouldSkipOnFailure) {
+                            // TODO: drop this in final patch
+                            LOGGER.error("Marking failed cluster {} as partial: {}", clusterAlias, e);
+                            computeListener.markAsPartial(clusterAlias, e);
+                            exchangeListener.onResponse(null);
+                        } else {
+                            exchangeListener.onFailure(e);
+                        }
                     })
                 );
             }