add failure handling and maxInFlighChunk mechanism

Author: drempapis

server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -45,11 +45,13 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.IntConsumer;
 import java.util.function.Supplier;
 
@@ -97,8 +99,18 @@ public void execute(
         SearchHits hits = null;
         try {
             // Collect all pending chunk futures
-            final List<CompletableFuture<Void>> pendingChunks = new ArrayList<>();
-            hits = buildSearchHits(context, docIdsToLoad, profiler, rankDocs, memoryChecker, writer, pendingChunks);
+            final int maxInFlightChunks = 1; // TODO make configurable
+            final ArrayDeque<CompletableFuture<Void>> pendingChunks = new ArrayDeque<>();
+            final AtomicReference<Throwable> sendFailure = new AtomicReference<>();
+            hits = buildSearchHits(context,
+                docIdsToLoad,
+                profiler,
+                rankDocs,
+                memoryChecker,
+                writer,
+                pendingChunks,
+                maxInFlightChunks,
+                sendFailure);
 
             // Wait for all chunks to be ACKed before setting final result
             if (writer != null && pendingChunks.isEmpty() == false) {
@@ -152,7 +164,7 @@ public void execute(SearchContext context, int[] docIdsToLoad, RankDocShardInfo
                 : Profilers.startProfilingFetchPhase();
         SearchHits hits = null;
         try {
-            hits = buildSearchHits(context, docIdsToLoad, profiler, rankDocs, memoryChecker, null, null);
+            hits = buildSearchHits(context, docIdsToLoad, profiler, rankDocs, memoryChecker, null, null, 0, null);
         } finally {
             try {
                 // Always finish profiling
@@ -187,7 +199,10 @@ private SearchHits buildSearchHits(
         RankDocShardInfo rankDocs,
         IntConsumer memoryChecker,
         FetchPhaseResponseChunk.Writer writer,
-        List<CompletableFuture<Void>> pendingChunks
+        ArrayDeque<CompletableFuture<Void>> pendingChunks,
+        int maxInFlightChunks,
+        AtomicReference<Throwable> sendFailure
+
     ) {
         var lookup = context.getSearchExecutionContext().getMappingLookup();
 
@@ -333,7 +348,9 @@ protected SearchHit nextDoc(int doc) throws IOException {
                 context.queryResult(),
                 writer,
                 5, // TODO set a proper number
-                pendingChunks
+                pendingChunks,
+                maxInFlightChunks,
+                sendFailure
             );
 
             if (context.isCancelled()) {

server/src/main/java/org/elasticsearch/search/fetch/FetchPhaseDocsIterator.java

@@ -25,11 +25,13 @@
 import org.elasticsearch.search.query.SearchTimeoutException;
 
 import java.io.IOException;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Given a set of doc ids and an index reader, sorts the docs by id (when not streaming),
@@ -110,7 +112,9 @@ public final IterateResult iterate(
         QuerySearchResult querySearchResult,
         FetchPhaseResponseChunk.Writer chunkWriter,
         int chunkSize,
-        List<CompletableFuture<Void>> pendingChunks
+        ArrayDeque<CompletableFuture<Void>> pendingChunks,
+        int maxInFlightChunks,
+        AtomicReference<Throwable> sendFailure
     ) {
         SearchHit[] searchHits = new SearchHit[docIds.length];
         DocIdToIndex[] docs = new DocIdToIndex[docIds.length];
@@ -172,26 +176,44 @@ public final IterateResult iterate(
                         if (chunkBuffer.isEmpty()) {
                             currentChunkSequenceStart = hitSequenceCounter.get();
                         }
-
                         // Assign sequence to this hit and increment counter
                         hitSequenceCounter.getAndIncrement();
 
                         chunkBuffer.add(hit);
 
                         // Send intermediate chunks - not when it's the last iteration
                         if (chunkBuffer.size() >= chunkSize && i < docs.length - 1) {
+                            // fail fast if any earlier send already failed
+                            Throwable knownFailure = sendFailure.get();
+                            if (knownFailure != null) {
+                                throw new RuntimeException("Fetch chunk failed", knownFailure);
+                            }
+
                             // Send chunk with sequence information
-                            pendingChunks.add(
-                                sendChunk(
+                            CompletableFuture<Void> chunkFuture = sendChunk(
                                     chunkWriter,
                                     chunkBuffer,
                                     shardId,
                                     currentChunkSequenceStart,
                                     i - chunkBuffer.size() + 1,
                                     docIds.length,
                                     Float.NaN
-                                )
-                            );
+                                );
+
+                            // record failures as soon as they happen
+                            chunkFuture.whenComplete((ok, ex) -> {
+                                if (ex != null) {
+                                    sendFailure.compareAndSet(null, ex);
+                                }
+                            });
+
+                            pendingChunks.addLast(chunkFuture);
+
+                            // Backpressure: bound in-flight
+                            if (pendingChunks.size() >= maxInFlightChunks) {
+                                awaitOldestOrFail(pendingChunks, sendFailure);
+                            }
+
                             chunkBuffer.clear();
                         }
                     } else {
@@ -246,6 +268,16 @@ public final IterateResult iterate(
         return new IterateResult(searchHits, lastChunk, lastChunkSequenceStart);
     }
 
+    private static void awaitOldestOrFail(ArrayDeque<CompletableFuture<Void>> inFlight, AtomicReference<Throwable> sendFailure) {
+        final CompletableFuture<Void> oldest = inFlight.removeFirst();
+        try {
+            oldest.get();
+        } catch (Exception e) {
+            sendFailure.compareAndSet(null, e);
+            throw new RuntimeException("Failed to send fetch chunk", e);
+        }
+    }
+
     /**
      * Sends a chunk of hits to the coordinator with sequence information for ordering.
      */

server/src/test/java/org/elasticsearch/search/fetch/FetchPhaseDocsIteratorTests.java

@@ -86,6 +86,8 @@ protected SearchHit nextDoc(int doc) {
             new QuerySearchResult(),
             null,
             0,
+            null,
+            0,
             null
         );
 
@@ -137,7 +139,7 @@ protected SearchHit nextDoc(int doc) {
 
         Exception e = expectThrows(
             FetchPhaseExecutionException.class,
-            () -> it.iterate(null, reader, docs, randomBoolean(), new QuerySearchResult(), null, 0, null)
+            () -> it.iterate(null, reader, docs, randomBoolean(), new QuerySearchResult(), null, 0, null, 0, null)
         );
         assertThat(e.getMessage(), containsString("Error running fetch phase for doc [" + badDoc + "]"));
         assertThat(e.getCause(), instanceOf(IllegalArgumentException.class));