Move shard failure after send failure into SegmentFileTransferHandler.

Signed-off-by: Marc Handalian <[email protected]>
opensearch-project · Jun 11, 2022 · ec397f1 · ec397f1
1 parent 2bfbf4a
commit ec397f1
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 30 deletions.
diff --git a/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java b/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java
@@ -142,7 +142,6 @@ public RecoverySourceHandler(
             logger,
             threadPool,
             cancellableThreads,
-            this::failEngine,
             fileChunkSizeInBytes,
             maxConcurrentFileChunks
         );
@@ -1050,8 +1049,4 @@ private void cleanFiles(
             }))
         );
     }
-
-    protected void failEngine(IOException cause) {
-        shard.failShard("recovery", cause);
-    }
 }
diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java b/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java
@@ -37,7 +37,6 @@
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.concurrent.ConcurrentLinkedDeque;
-import java.util.function.Consumer;
 import java.util.function.IntSupplier;
 
 /**
@@ -57,7 +56,6 @@ public final class SegmentFileTransferHandler {
     private final int maxConcurrentFileChunks;
     private final DiscoveryNode targetNode;
     private final CancellableThreads cancellableThreads;
-    private final Consumer<IOException> onCorruptException;
 
     public SegmentFileTransferHandler(
         IndexShard shard,
@@ -66,7 +64,6 @@ public SegmentFileTransferHandler(
         Logger logger,
         ThreadPool threadPool,
         CancellableThreads cancellableThreads,
-        Consumer<IOException> onCorruptException,
         int fileChunkSizeInBytes,
         int maxConcurrentFileChunks
     ) {
@@ -76,7 +73,6 @@ public SegmentFileTransferHandler(
         this.logger = logger;
         this.threadPool = threadPool;
         this.cancellableThreads = cancellableThreads;
-        this.onCorruptException = onCorruptException;
         this.chunkSizeInBytes = fileChunkSizeInBytes;
         // if the target is on an old version, it won't be able to handle out-of-order file chunks.
         this.maxConcurrentFileChunks = maxConcurrentFileChunks;
@@ -184,7 +180,7 @@ public void handleErrorOnSendFiles(Store store, Exception e, StoreFileMetadata[]
                     if (localException == null) {
                         localException = corruptIndexException;
                     }
-                    onCorruptException.accept(corruptIndexException);
+                    shard.failShard("error sending files", corruptIndexException);
                 }
             }
             if (localException != null) {

diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationSourceHandler.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationSourceHandler.java
@@ -30,7 +30,6 @@
 import org.opensearch.transport.Transports;
 
 import java.io.Closeable;
-import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
@@ -84,7 +83,6 @@ class SegmentReplicationSourceHandler {
             logger,
             threadPool,
             cancellableThreads,
-            this::failEngine,
             fileChunkSizeInBytes,
             maxConcurrentFileChunks
         );
@@ -156,10 +154,6 @@ public void cancel(String reason) {
         cancellableThreads.cancel(reason);
     }
 
-    private void failEngine(IOException e) {
-        shard.failShard("Failed Replication", e);
-    }
-
     CopyState getCopyState() {
         return copyState;
     }

diff --git a/server/src/test/java/org/opensearch/indices/recovery/RecoverySourceHandlerTests.java b/server/src/test/java/org/opensearch/indices/recovery/RecoverySourceHandlerTests.java
@@ -547,6 +547,11 @@ public void writeFileChunk(
         };
         IndexShard mockShard = mock(IndexShard.class);
         when(mockShard.shardId()).thenReturn(new ShardId("testIndex", "testUUID", 0));
+        doAnswer(invocation -> {
+            assertFalse(failedEngine.get());
+            failedEngine.set(true);
+            return null;
+        }).when(mockShard).failShard(any(), any());
         RecoverySourceHandler handler = new RecoverySourceHandler(
             mockShard,
             new AsyncRecoveryTarget(target, recoveryExecutor),
@@ -555,13 +560,7 @@ public void writeFileChunk(
             Math.toIntExact(recoverySettings.getChunkSize().getBytes()),
             between(1, 8),
             between(1, 8)
-        ) {
-            @Override
-            protected void failEngine(IOException cause) {
-                assertFalse(failedEngine.get());
-                failedEngine.set(true);
-            }
-        };
+        );
         SetOnce<Exception> sendFilesError = new SetOnce<>();
         CountDownLatch latch = new CountDownLatch(1);
         handler.sendFiles(
@@ -623,6 +622,11 @@ public void writeFileChunk(
         };
         IndexShard mockShard = mock(IndexShard.class);
         when(mockShard.shardId()).thenReturn(new ShardId("testIndex", "testUUID", 0));
+        doAnswer(invocation -> {
+            assertFalse(failedEngine.get());
+            failedEngine.set(true);
+            return null;
+        }).when(mockShard).failShard(any(), any());
         RecoverySourceHandler handler = new RecoverySourceHandler(
             mockShard,
             new AsyncRecoveryTarget(target, recoveryExecutor),
@@ -631,13 +635,7 @@ public void writeFileChunk(
             Math.toIntExact(recoverySettings.getChunkSize().getBytes()),
             between(1, 10),
             between(1, 4)
-        ) {
-            @Override
-            protected void failEngine(IOException cause) {
-                assertFalse(failedEngine.get());
-                failedEngine.set(true);
-            }
-        };
+        );
         PlainActionFuture<Void> sendFilesFuture = new PlainActionFuture<>();
         handler.sendFiles(store, metas.toArray(new StoreFileMetadata[0]), () -> 0, sendFilesFuture);
         Exception ex = expectThrows(Exception.class, sendFilesFuture::actionGet);