Make peer recovery send file chunks async

server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java

@@ -39,6 +39,7 @@
 import org.elasticsearch.common.StopWatch;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.logging.Loggers;
@@ -81,6 +82,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.IntSupplier;
 import java.util.stream.StreamSupport;
@@ -704,12 +706,12 @@ private class MultiFileSender extends ActionRunnable<Void> implements Closeable
         private final IntSupplier translogOps;
         private final LocalCheckpointTracker requestSeqIdTracker = new LocalCheckpointTracker(NO_OPS_PERFORMED, NO_OPS_PERFORMED);
         private final Semaphore semaphore = new Semaphore(0);
-        private boolean closed = false;
         private final Iterator<StoreFileMetaData> remainingFiles;
         private StoreFileMetaData currentFile;
         private InputStreamIndexInput currentInput = null;
         private long currentChunkPosition = 0;
         private final Deque<byte[]> recycledBuffers = ConcurrentCollections.newDeque();
+        private final AtomicReference<Tuple<StoreFileMetaData, Exception>> error = new AtomicReference<>();
 
         MultiFileSender(Store store, IntSupplier translogOps, StoreFileMetaData[] files, ActionListener<Void> listener) {
             super(ActionListener.notifyOnce(listener));
@@ -724,6 +726,10 @@ protected void doRun() throws Exception {
             assert Transports.assertNotTransportThread(RecoverySourceHandler.this + "[send file chunk]");
             while (true) {
                 assert semaphore.availablePermits() == 0;
+                if (error.get() != null) {
+                    handleErrorOnSendFiles(store, error.get().v2(), new StoreFileMetaData[]{error.get().v1()});
+                    throw error.get().v2();
+                }
                 final FileChunk chunk = readNextChunk();
                 if (chunk == null) {
                     semaphore.release(); // allow other threads respond if we are not done yet.
@@ -743,25 +749,20 @@ protected void doRun() throws Exception {
                                     sendFileExecutor.execute(this); // fork off from the network thread
                                 }
                             },
-                            // need to fork as `handleErrorOnSendFiles` might read some files which should not happen on the network thread
-                            e -> sendFileExecutor.execute(new ActionRunnable<>(this.listener) {
-                                @Override
-                                protected void doRun() throws Exception {
-                                    cancellableThreads.execute(semaphore::acquire);
-                                    try (Releasable ignored = semaphore::release) {
-                                        handleErrorOnSendFiles(store, e, new StoreFileMetaData[]{chunk.md});
-                                        throw e;
-                                    }
+                            e -> {
+                                if (error.compareAndSet(null, Tuple.tuple(chunk.md, e)) && semaphore.tryAcquire()) {
+                                    // have to fork as handleErrorOnSendFiles can read file which should not happen on the network thread.
+                                    sendFileExecutor.execute(this);
                                 }
                             })
-                        )
                     )
                 );
                 if (canSendMore() == false) {
                     semaphore.release();
                     // Here we have to retry before abort to avoid a race situation where the other threads have flipped `canSendMore`
                     // condition but they are not going to resume the sending process because this thread still holds the semaphore.
-                    if (canSendMore() == false || semaphore.tryAcquire() == false) {
+                    final boolean changed = canSendMore() || error.get() != null;
+                    if (changed == false || semaphore.tryAcquire() == false) {
                         break;
                     }
                 }
@@ -770,9 +771,6 @@ protected void doRun() throws Exception {
 
         FileChunk readNextChunk() throws Exception {
             assert semaphore.availablePermits() == 0;
-            if (closed) {
-                throw new IllegalStateException("MultiFileSender was closed");
-            }
             try {
                 if (currentInput == null) {
                     if (remainingFiles.hasNext() == false) {
@@ -816,10 +814,7 @@ boolean canSendMore() {
         @Override
         public void close() throws IOException {
             assert semaphore.availablePermits() == 0;
-            if (closed == false) {
-                closed = true;
-                IOUtils.close(recycledBuffers::clear, currentInput, () -> currentInput = null);
-            }
+            IOUtils.close(recycledBuffers::clear, currentInput, () -> currentInput = null);
         }
     }
 

server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java

@@ -107,7 +107,6 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
-import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.core.IsNull.notNullValue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyBoolean;
@@ -673,8 +672,7 @@ public void writeFileChunk(StoreFileMetaData md, long position, BytesReference c
             assertThat(error.get(), notNullValue());
             assertThat(error.get().getMessage(), containsString("test chunk exception"));
         });
-        assertThat("no more chunks should be sent", sentChunks.get(),
-            lessThanOrEqualTo(Math.min(totalChunks, maxConcurrentChunks * 2)));
+        assertThat("no more chunks should be sent", sentChunks.get(), equalTo(Math.min(totalChunks, maxConcurrentChunks)));
         store.close();
     }