Retry on shard failure

Author: dnhatn

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

@@ -167,7 +167,6 @@ static TransportVersion def(int id) {
     public static final TransportVersion RANK_DOC_OPTIONAL_METADATA_FOR_EXPLAIN = def(8_833_00_0);
     public static final TransportVersion ESQL_RETRY_ON_SHARD_LEVEL_FAILURE = def(8_834_00_0);
 
-
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

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

@@ -45,19 +45,23 @@ public FailureCollector(int maxExceptions) {
         this.nonCancelledExceptionsPermits = new Semaphore(maxExceptions);
     }
 
-    private static Exception unwrapTransportException(TransportException te) {
-        final Throwable cause = te.getCause();
-        if (cause == null) {
-            return te;
-        } else if (cause instanceof Exception ex) {
-            return ex;
+    public static Exception unwrapTransportException(Exception e) {
+        if (e instanceof TransportException te) {
+            final Throwable cause = te.getCause();
+            if (cause == null) {
+                return te;
+            } else if (cause instanceof Exception ex) {
+                return ex;
+            } else {
+                return new ElasticsearchException(cause);
+            }
         } else {
-            return new ElasticsearchException(cause);
+            return e;
         }
     }
 
     public void unwrapAndCollect(Exception e) {
-        e = e instanceof TransportException te ? unwrapTransportException(te) : e;
+        e = unwrapTransportException(e);
         if (ExceptionsHelper.unwrap(e, TaskCancelledException.class) != null) {
             if (nonCancelledExceptions.isEmpty() && cancelledExceptionsPermits.tryAcquire()) {
                 cancelledExceptions.add(e);

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

@@ -65,6 +65,10 @@ public Page takePage() {
         return page;
     }
 
+    public boolean hasData() {
+        return page != null;
+    }
+
     public long ramBytesUsedByPage() {
         if (page != null) {
             return page.ramBytesUsedByBlocks();

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

@@ -51,7 +51,8 @@
 /**
  * {@link ExchangeService} is responsible for exchanging pages between exchange sinks and sources on the same or different nodes.
  * It holds a map of {@link ExchangeSinkHandler} instances for each node in the cluster to serve {@link ExchangeRequest}s
- * To connect exchange sources to exchange sinks, use {@link ExchangeSourceHandler#addRemoteSink(RemoteSink, boolean, int, ActionListener)}.
+ * To connect exchange sources to exchange sinks,
+ * use {@link ExchangeSourceHandler#addRemoteSink(RemoteSink, boolean, Runnable, int, ActionListener)}.
  */
 public final class ExchangeService extends AbstractLifecycleComponent {
     // TODO: Make this a child action of the data node transport to ensure that exchanges

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

@@ -22,10 +22,10 @@
 
 /**
  * An {@link ExchangeSinkHandler} receives pages and status from its {@link ExchangeSink}s, which are created using
- * {@link #createExchangeSink()}} method. Pages and status can then be retrieved asynchronously by {@link ExchangeSourceHandler}s
+ * {@link #createExchangeSink(Runnable)}} method. Pages and status can then be retrieved asynchronously by {@link ExchangeSourceHandler}s
  * using the {@link #fetchPageAsync(boolean, ActionListener)} method.
  *
- * @see #createExchangeSink()
+ * @see #createExchangeSink(Runnable)
  * @see #fetchPageAsync(boolean, ActionListener)
  * @see ExchangeSourceHandler
  */
@@ -53,15 +53,18 @@ public ExchangeSinkHandler(BlockFactory blockFactory, int maxBufferSize, LongSup
     private class ExchangeSinkImpl implements ExchangeSink {
         boolean finished;
         private final SubscribableListener<Void> onFinished = new SubscribableListener<>();
+        private final Runnable onPageAdded;
 
-        ExchangeSinkImpl() {
+        ExchangeSinkImpl(Runnable onPageAdded) {
             onChanged();
+            this.onPageAdded = onPageAdded;
             buffer.addCompletionListener(onFinished);
             outstandingSinks.incrementAndGet();
         }
 
         @Override
         public void addPage(Page page) {
+            onPageAdded.run();
             buffer.addPage(page);
             notifyListeners();
         }
@@ -101,7 +104,7 @@ public IsBlockedResult waitForWriting() {
      * @param sourceFinished if true, then this handler can finish as sources have enough pages.
      * @param listener       the listener that will be notified when pages are ready or this handler is finished
      * @see RemoteSink
-     * @see ExchangeSourceHandler#addRemoteSink(RemoteSink, boolean, int, ActionListener)
+     * @see ExchangeSourceHandler#addRemoteSink(RemoteSink, boolean, Runnable, int, ActionListener)
      */
     public void fetchPageAsync(boolean sourceFinished, ActionListener<ExchangeResponse> listener) {
         if (sourceFinished) {
@@ -161,10 +164,11 @@ private void notifyListeners() {
     /**
      * Create a new exchange sink for exchanging data
      *
+     * @param onPageAdded a {@link Runnable} that called when a new page is added
      * @see ExchangeSinkOperator
      */
-    public ExchangeSink createExchangeSink() {
-        return new ExchangeSinkImpl();
+    public ExchangeSink createExchangeSink(Runnable onPageAdded) {
+        return new ExchangeSinkImpl(onPageAdded);
     }
 
     /**

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

@@ -27,10 +27,10 @@
 /**
  * An {@link ExchangeSourceHandler} asynchronously fetches pages and status from multiple {@link RemoteSink}s
  * and feeds them to its {@link ExchangeSource}, which are created using the {@link #createExchangeSource()}) method.
- * {@link RemoteSink}s are added using the {@link #addRemoteSink(RemoteSink, boolean, int, ActionListener)}) method.
+ * {@link RemoteSink}s are added using the {@link #addRemoteSink(RemoteSink, boolean, Runnable, int, ActionListener)}) method.
  *
  * @see #createExchangeSource()
- * @see #addRemoteSink(RemoteSink, boolean, int, ActionListener)
+ * @see #addRemoteSink(RemoteSink, boolean, Runnable, int, ActionListener)
  */
 public final class ExchangeSourceHandler {
     private final ExchangeBuffer buffer;
@@ -186,11 +186,13 @@ private final class RemoteSinkFetcher {
         private final RemoteSink remoteSink;
         private final boolean failFast;
         private final ActionListener<Void> completionListener;
+        private final Runnable onPageFetched;
 
-        RemoteSinkFetcher(RemoteSink remoteSink, boolean failFast, ActionListener<Void> completionListener) {
+        RemoteSinkFetcher(RemoteSink remoteSink, boolean failFast, Runnable onPageFetched, ActionListener<Void> completionListener) {
             outstandingSinks.trackNewInstance();
             this.remoteSink = remoteSink;
             this.failFast = failFast;
+            this.onPageFetched = onPageFetched;
             this.completionListener = completionListener;
         }
 
@@ -203,6 +205,7 @@ void fetchPage() {
                 remoteSink.fetchPageAsync(toFinishSinks, ActionListener.wrap(resp -> {
                     Page page = resp.takePage();
                     if (page != null) {
+                        onPageFetched.run();
                         buffer.addPage(page);
                     }
                     if (resp.finished()) {
@@ -259,12 +262,19 @@ void onSinkComplete() {
      *                   - If {@code true}, failures from this remote sink will cause the exchange source to abort.
      *                   Callers can safely ignore failures notified via this listener, as they are collected and
      *                   reported by the exchange source.
+     * @param onPageFetched a callback that is called when a new page is fetched from the remote sink
      * @param instances  the number of concurrent ``clients`` that this handler should use to fetch pages.
      *                   More clients reduce latency, but add overhead.
      * @param listener   a listener that will be notified when the sink fails or completes
      * @see ExchangeSinkHandler#fetchPageAsync(boolean, ActionListener)
      */
-    public void addRemoteSink(RemoteSink remoteSink, boolean failFast, int instances, ActionListener<Void> listener) {
+    public void addRemoteSink(
+        RemoteSink remoteSink,
+        boolean failFast,
+        Runnable onPageFetched,
+        int instances,
+        ActionListener<Void> listener
+    ) {
         final int sinkId = nextSinkId.incrementAndGet();
         remoteSinks.put(sinkId, remoteSink);
         final ActionListener<Void> sinkListener = ActionListener.assertAtLeastOnce(
@@ -284,7 +294,7 @@ public void onFailure(Exception e) {
             protected void doRun() {
                 try (EsqlRefCountingListener refs = new EsqlRefCountingListener(sinkListener)) {
                     for (int i = 0; i < instances; i++) {
-                        var fetcher = new RemoteSinkFetcher(remoteSink, failFast, refs.acquire());
+                        var fetcher = new RemoteSinkFetcher(remoteSink, failFast, onPageFetched, refs.acquire());
                         fetcher.fetchPage();
                     }
                 }

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

@@ -297,7 +297,7 @@ public void testEarlyTermination() {
             final int maxAllowedRows = between(1, 100);
             final AtomicInteger processedRows = new AtomicInteger(0);
             var sinkHandler = new ExchangeSinkHandler(driverContext.blockFactory(), positions, System::currentTimeMillis);
-            var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(), Function.identity());
+            var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(() -> {}), Function.identity());
             final var delayOperator = new EvalOperator(driverContext.blockFactory(), new EvalOperator.ExpressionEvaluator() {
                 @Override
                 public Block eval(Page page) {
@@ -335,7 +335,7 @@ public void testResumeOnEarlyFinish() throws Exception {
             var sourceHandler = new ExchangeSourceHandler(between(1, 5), threadPool.executor("esql"), sourceFuture);
             var sinkHandler = new ExchangeSinkHandler(driverContext.blockFactory(), between(1, 5), System::currentTimeMillis);
             var sourceOperator = new ExchangeSourceOperator(sourceHandler.createExchangeSource());
-            var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(), Function.identity());
+            var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(() -> {}), Function.identity());
             Driver driver = new Driver(driverContext, sourceOperator, List.of(), sinkOperator, () -> {});
             PlainActionFuture<Void> future = new PlainActionFuture<>();
             Driver.start(threadPool.getThreadContext(), threadPool.executor("esql"), driver, between(1, 1000), future);

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

@@ -220,6 +220,7 @@ List<Driver> createDriversForInput(List<Page> input, List<Page> results, boolean
         sourceExchanger.addRemoteSink(
             sinkExchanger::fetchPageAsync,
             randomBoolean(),
+            () -> {},
             1,
             ActionListener.<Void>noop().delegateResponse((l, e) -> {
                 throw new AssertionError("unexpected failure", e);
@@ -248,7 +249,7 @@ List<Driver> createDriversForInput(List<Page> input, List<Page> results, boolean
                         simpleWithMode(AggregatorMode.INTERMEDIATE).get(driver1Context),
                         intermediateOperatorItr.next()
                     ),
-                    new ExchangeSinkOperator(sinkExchanger.createExchangeSink(), Function.identity()),
+                    new ExchangeSinkOperator(sinkExchanger.createExchangeSink(() -> {}), Function.identity()),
                     () -> {}
                 )
             );

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

@@ -97,12 +97,12 @@ public void testBasic() throws Exception {
             pages[i] = new Page(blockFactory.newConstantIntBlockWith(i, 2));
         }
         ExchangeSinkHandler sinkExchanger = new ExchangeSinkHandler(blockFactory, 2, threadPool.relativeTimeInMillisSupplier());
-        ExchangeSink sink1 = sinkExchanger.createExchangeSink();
-        ExchangeSink sink2 = sinkExchanger.createExchangeSink();
+        ExchangeSink sink1 = sinkExchanger.createExchangeSink(() -> {});
+        ExchangeSink sink2 = sinkExchanger.createExchangeSink(() -> {});
         PlainActionFuture<Void> sourceCompletion = new PlainActionFuture<>();
         ExchangeSourceHandler sourceExchanger = new ExchangeSourceHandler(3, threadPool.executor(ESQL_TEST_EXECUTOR), sourceCompletion);
         ExchangeSource source = sourceExchanger.createExchangeSource();
-        sourceExchanger.addRemoteSink(sinkExchanger::fetchPageAsync, randomBoolean(), 1, ActionListener.noop());
+        sourceExchanger.addRemoteSink(sinkExchanger::fetchPageAsync, randomBoolean(), () -> {}, 1, ActionListener.noop());
         SubscribableListener<Void> waitForReading = source.waitForReading().listener();
         assertFalse(waitForReading.isDone());
         assertNull(source.pollPage());
@@ -340,10 +340,16 @@ public void testConcurrentWithHandlers() {
                 sinkHandler = randomFrom(sinkHandlers);
             } else {
                 sinkHandler = new ExchangeSinkHandler(blockFactory, randomExchangeBuffer(), threadPool.relativeTimeInMillisSupplier());
-                sourceExchanger.addRemoteSink(sinkHandler::fetchPageAsync, randomBoolean(), randomIntBetween(1, 3), ActionListener.noop());
+                sourceExchanger.addRemoteSink(
+                    sinkHandler::fetchPageAsync,
+                    randomBoolean(),
+                    () -> {},
+                    randomIntBetween(1, 3),
+                    ActionListener.noop()
+                );
                 sinkHandlers.add(sinkHandler);
             }
-            return sinkHandler.createExchangeSink();
+            return sinkHandler.createExchangeSink(() -> {});
         };
         final int maxInputSeqNo = rarely() ? -1 : randomIntBetween(0, 50_000);
         final int maxOutputSeqNo = rarely() ? -1 : randomIntBetween(0, 50_000);
@@ -398,14 +404,14 @@ public void testExchangeSourceContinueOnFailure() {
                         l.onResponse(new ExchangeResponse(blockFactory, page, r.finished()));
                     }));
                 }
-            }, false, instance, ActionListener.wrap(r -> {
+            }, false, () -> {}, instance, ActionListener.wrap(r -> {
                 assertFalse(sinkFailed.get());
                 completedSinks.incrementAndGet();
             }, e -> {
                 assertTrue(sinkFailed.get());
                 failedSinks.incrementAndGet();
             }));
-            return sinkHandler.createExchangeSink();
+            return sinkHandler.createExchangeSink(() -> {});
         };
         Set<Integer> actualSeqNos = runConcurrentTest(
             maxInputSeqNo,
@@ -430,7 +436,7 @@ public void testClosingSinks() {
         Page p1 = new Page(block1);
         Page p2 = new Page(block2);
         ExchangeSinkHandler sinkExchanger = new ExchangeSinkHandler(blockFactory, 2, threadPool.relativeTimeInMillisSupplier());
-        ExchangeSink sink = sinkExchanger.createExchangeSink();
+        ExchangeSink sink = sinkExchanger.createExchangeSink(() -> {});
         sink.addPage(p1);
         sink.addPage(p2);
         assertFalse(sink.waitForWriting().listener().isDone());
@@ -475,7 +481,7 @@ public void testFinishEarly() throws Exception {
                     throw new AssertionError(e);
                 }
             }
-        }, false, between(1, 3), sinkCompleted);
+        }, false, () -> {}, between(1, 3), sinkCompleted);
         threadPool.schedule(
             () -> sourceHandler.finishEarly(randomBoolean(), ActionListener.noop()),
             TimeValue.timeValueMillis(between(0, 10)),
@@ -526,6 +532,7 @@ public void testConcurrentWithTransportActions() {
             sourceHandler.addRemoteSink(
                 exchange0.newRemoteSink(task, exchangeId, node0, connection),
                 randomBoolean(),
+                () -> {},
                 randomIntBetween(1, 5),
                 ActionListener.noop()
             );
@@ -535,7 +542,7 @@ public void testConcurrentWithTransportActions() {
                 maxInputSeqNo,
                 maxOutputSeqNo,
                 sourceHandler::createExchangeSource,
-                sinkHandler::createExchangeSink
+                () -> sinkHandler.createExchangeSink(() -> {})
             );
             var expectedSeqNos = IntStream.range(0, Math.min(maxInputSeqNo, maxOutputSeqNo)).boxed().collect(Collectors.toSet());
             assertThat(actualSeqNos, hasSize(expectedSeqNos.size()));
@@ -601,12 +608,18 @@ public void sendResponse(TransportResponse transportResponse) {
             sourceHandler.addRemoteSink(
                 exchange0.newRemoteSink(task, exchangeId, node0, connection),
                 true,
+                () -> {},
                 randomIntBetween(1, 5),
                 ActionListener.noop()
             );
             Exception err = expectThrows(
                 Exception.class,
-                () -> runConcurrentTest(maxSeqNo, maxSeqNo, sourceHandler::createExchangeSource, sinkHandler::createExchangeSink)
+                () -> runConcurrentTest(
+                    maxSeqNo,
+                    maxSeqNo,
+                    sourceHandler::createExchangeSource,
+                    () -> sinkHandler.createExchangeSink(() -> {})
+                )
             );
             Throwable cause = ExceptionsHelper.unwrap(err, IOException.class);
             assertNotNull(cause);